agent_brain.py

# web_agent/agent_brain.py
import sys
import os
import json

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import ai_interpret
from config import WEBSITE_URLS_STR

from pydantic import BaseModel, ConfigDict, Field
from typing import Optional, Union, Dict, Any
import json
import re

try:
    from langfuse.decorators import observe
except ImportError:
    # Fallback if langfuse isn't installed locally yet
    def observe(*args, **kwargs):
        def decorator(func):
            return func
        return decorator

class AgentActionSchema(BaseModel):
    model_config = ConfigDict(extra="forbid")
    
    action: str
    thought: str = Field(default="")
    wait_time: Optional[int] = Field(default=2)
    element_id: Optional[int] = Field(default=None)
    text: Optional[str] = Field(default=None)
    value: Optional[str] = Field(default=None)
    key: Optional[str] = Field(default=None)
    url: Optional[str] = Field(default=None)
    selector: Optional[str] = Field(default=None)

@observe()
def _extract_json(text: str) -> dict:
    if not text or not isinstance(text, str):
        return {"action": "wait", "wait_time": 5, "thought": "Received empty or invalid text from LLM."}

    text = text.replace("```json", "").replace("```", "").strip()
    match = re.search(r'\{.*\}', text, re.DOTALL)
    
    if not match:
        print(f"⚠️ [JSON ERROR] No JSON object found in response.")
        return {"action": "wait", "wait_time": 5, "thought": "No JSON found in response."}
        
    json_str = match.group(0).replace("<|im_end|>", "").strip()
    
    try:
        # Pydantic Strict Validation (extra='forbid')
        action_obj = AgentActionSchema.model_validate_json(json_str)
        return action_obj.model_dump(exclude_none=True)
    except Exception as e:
        print(f"⚠️ [SCHEMA VALIDATION ERROR] {str(e)}")
        # Fallback to wait if validation fails (silent leak prevented)
        return {"action": "wait", "wait_time": 5, "thought": f"Schema Validation Failed: {str(e)}"}


def _extract_clean_json(text: str) -> dict:
    return _extract_json(text)

def _is_failed_action(action_data: dict) -> bool:
    if not action_data or not isinstance(action_data, dict):
        return True
    thought = action_data.get("thought", "")
    if "Failed to parse JSON" in thought or "No JSON found in response" in thought or "Parsed result is not a dictionary" in thought:
        return True
    if "action" not in action_data:
        return True
    return False

def extract_goal_and_url(user_prompt: str) -> dict:
    print("\n🧠 Target URL aur Goal samajh raha hai...")
    messages = [
        {
            "role": "system",
            "content": f"""Extract the target website URL, the core goal, the target Facebook page name (if mentioned in the prompt), and the structured payload from the user's prompt.
KNOWN WEBSITES:
{WEBSITE_URLS_STR}

OUTPUT STRICT JSON FORMAT EXACTLY LIKE THIS:
{{
  "start_url": "https://www.linkedin.com/", 
  "goal": "Search for X, click their profile, and connect.",
  "page_name": "TrendFollower",
  "payload": {{
     "image_path": null,
     "title": "Exact Title Here",
     "description": "Exact Description Here",
     "link": "https://link.com"
  }}
}}
If target Facebook page name is not mentioned, set "page_name" to null. If any payload field is not mentioned, set it to null.
Do NOT write any introduction or explanation. Do NOT write markdown blocks. Just output raw JSON."""
        },
        {"role": "user", "content": f"Task: {user_prompt}"}
    ]
    
    parsed = {}
    for attempt_idx in range(1, 4):
        try:
            print(f"🧠 [EXTRACT GOAL] LLM Call Attempt {attempt_idx}/3...")
            raw = ai_interpret.get_gemma_response_with_fallback(messages, max_new_tokens=1500, temperature=0.1)
            parsed = _extract_clean_json(raw)
            if parsed and ("goal" in parsed or "start_url" in parsed):
                break
            else:
                print(f"⚠️ [EXTRACT GOAL] Invalid response format on attempt {attempt_idx}. Retrying...")
        except Exception as e:
            print(f"⚠️ [EXTRACT GOAL] Exception on attempt {attempt_idx}: {e}")
            parsed = {}
            
    if not parsed:
        parsed = {
            "start_url": "https://www.google.com/",
            "goal": user_prompt,
            "page_name": None,
            "payload": {}
        }
        
    # 🚀 Rule-based Start URL Fallback Safety Net (in case LLM fails or is rate-limited)
    start_url = parsed.get("start_url") or ""
    if not start_url or "google.com" in start_url:
        prompt_lower = user_prompt.lower()
        if "linkedin" in prompt_lower:
            parsed["start_url"] = "https://www.linkedin.com/"
        elif "pinterest" in prompt_lower:
            parsed["start_url"] = "https://www.pinterest.com/"
        elif "facebook" in prompt_lower:
            parsed["start_url"] = "https://www.facebook.com/"
        elif "instagram" in prompt_lower:
            parsed["start_url"] = "https://www.instagram.com/"
        elif "tumblr" in prompt_lower:
            parsed["start_url"] = "https://www.tumblr.com/"
        elif "medium" in prompt_lower:
            parsed["start_url"] = "https://medium.com/"
            
    return parsed

def decide_action(goal, page_snapshot, action_history, attempt, max_attempts, image_path=None, image_b64=None, orig_w=1000, orig_h=1000, force_lightning_ai=False) -> dict:
    action_data = None
    history_str = ""
    if action_history:
        last_actions = action_history[-6:]
        history_str = "RECENT ACTIONS (PAY ATTENTION TO WHAT FAILED):\n"
        for a in last_actions:
            if a.get("_failed"):
                status = f"❌ {a.get('_error', 'Failed')}"
            else:
                status = "✅ Executed"
            history_str += f"  {status} -> {a.get('action')} | ID: {a.get('element_id', 'N/A')} | {a.get('text', a.get('file_path', a.get('url', '')))}\n"
    
    progress_str = "\n📊 GOAL PROGRESS (WHAT IS ALREADY DONE. DO NOT REPEAT THESE):\n"
    successful_actions = [a for a in action_history if not a.get("_failed")]
    
    uploads_done = [a for a in successful_actions if a.get("action") == "upload_file"]
    if uploads_done:
        progress_str += f"  ✅ Image/File Uploaded Successfully!\n"
        
    types_done = [a for a in successful_actions if a.get("action") in ["type", "type_and_enter"]]
    if types_done:
        for t in types_done:
            txt = str(t.get('text', ''))[:30]
            progress_str += f"  ✅ Typed: \"{txt}\"\n"

    # --- STANDARD DOM PATH (Gemma, Qwen, Groq, HF) ---
    messages = [
        {
            "role": "system",
            "content": f"""[ROLE]
You are a highly intelligent, AUTONOMOUS Web Agent. You act like a PATIENT HUMAN. 
You are also provided with a live screenshot of the current page state. Look at the screenshot to verify if fields are empty or filled, check for errors or popups, and locate elements visually before deciding.

[AVAILABLE_ACTIONS]
click, click_coordinates, type, type_and_enter, hover, scroll_down, scroll_up, go_back, refresh_page, go_to_url, press_key, upload_file, wait, task_done

[PROCESS_FLOW & DEEP REASONING]

RULE 1: STRICT DYNAMIC WAIT TIME -> CRITICAL
- Every action takes time to process. You MUST include a "wait_time" (in seconds).
- NEVER EXCEED 10 SECONDS.
   * Use 1 to 2 for simple UI changes (dropdowns, typing).
   * Use 3 to 10 for heavy tasks (uploading, clicking Publish, opening dialogs).

RULE 2: ABSOLUTE URL COMPLIANCE & ZERO HALLUCINATION -> CRITICAL
- You are STRICTLY FORBIDDEN from guessing, inventing, or predicting any URLs or page paths (e.g., NEVER invent paths like '/create', '/new', or '/post').
- When using the 'go_to_url' action, the URL MUST EXACTLY MATCH the 'start_url' or the explicit 'link' provided to you in the user's GOAL. 
- If the GOAL does not explicitly command you to navigate to a new specific URL, you are NOT ALLOWED to use 'go_to_url'. 
- Instead, rely COMPLETELY on interacting with the current PAGE STATE using 'click', 'type', or 'scroll_down' to find your way forward. Trust the user's entry point 100%.

RULE 3: THE POPUP AWARENESS 
- If a new dialog box appears, look for elements labeled '(INSIDE POPUP)'. Do NOT click background IDs.

RULE 4: NEVER REPEAT SAME CLICK & MOVE FORWARD -> ABSOLUTELY CRITICAL
- YOU MUST NEVER click the exact same button or input element consecutively, and NEVER try to reopen a dialog/composer that is already open.
- If 'RECENT ACTIONS' shows you already successfully clicked to open a dialog (e.g., clicking the 'What's on your mind' box), and the current DOM snapshot has elements with '(INSIDE POPUP)' or shows the composer is active, then the dialog is ALREADY open!
- YOU ARE STRICTLY FORBIDDEN from clicking to open it again or repeating that click. YOU MUST MOVE FORWARD TO THE NEXT STEP immediately (e.g., uploading the file or typing the content) instead of repeating successful past actions.

RULE 5: ANTI-SPAM & SINGLE-FIELD TYPING -> EXTREMELY CRITICAL
- Look closely at 'RECENT ACTIONS' and 'GOAL PROGRESS'.
- If 'RECENT ACTIONS' shows that you have already successfully executed a 'type' action on a field (e.g., typing description or title), you are STRICTLY FORBIDDEN from typing into that element ID again!
- If you believe a field is still empty, verify the DOM element ID. If you already typed into it once, assume it is filled and immediately proceed to type into the NEXT empty input field (e.g., if title is filled, type the description, then type the link).
- NEVER type into the same element ID consecutively!

RULE 6: EXIT RULE (TASK DONE) -> CRITICAL
- If your goal is complete based on the new PAGE STATE's Title, URL, or Success Text (e.g., "Saved", "Published", "Pending"), you MUST stop.
- Output EXACTLY: {{"action": "task_done"}}

- RULE 7: IMAGE/FILE UPLOADS AND OS POPUPS -> ABSOLUTELY CRITICAL
- YOU ARE STRICTLY FORBIDDEN from clicking on image/file upload trigger elements (e.g., buttons/containers saying "Add photos or videos", photo/video icons, or file drag-and-drop containers) using the 'click' or 'click_coordinates' action!
- REASON: Clicking these elements opens a native OS file dialog that blocks browser automation and causes the bot to fail.
- CORRECT ACTION: You MUST directly use the 'upload_file' action on the target upload element/container ID (e.g. the "Add photos or videos" container). This performs deep image/file URL injection directly into the target input element. The backend script will automatically locate the hidden file input inside it and inject the image path safely without triggering the blocking OS dialog.
- Facebook Main Feed Page: You MUST click ONLY on the text box saying "What's on your mind..." (e.g. "What's on your mind, TrendFollower?") to open the composer modal. YOU ARE STRICTLY FORBIDDEN from clicking the "Photo/video" button or icon on the main timeline page feed! Clicking that button immediately triggers the blocking Windows OS file dialog. Always open the composer modal first by clicking the "What's on your mind" text box.
- Facebook Page Composer: Once the post description is typed and 'upload_file' has successfully run on the photo/video container ID, locate the 'Post' button's ID (at the bottom of the composer popup) and click it to publish.

RULE 8: COORDINATE CLICK RULE -> CRITICAL
- If there is a blocking popup, popup close 'X' button, tour/tutorial overlay (like "Why stop at one?"), or any un-clickable element that is preventing you from proceeding, and it DOES NOT have an 'ai-id' in the PAGE STATE or if standard clicking fails, you can click it visually by coordinates.
- The browser window coordinates range from x = 0 (left-most) to 1280 (right-most), and y = 0 (top-most) to 800 (bottom-most).
- Look closely at the screenshot, identify where the element/button/icon is located, estimate the (x, y) coordinates on the 1280x800 screen, and perform a coordinate click to close or click it.

RULE 9: SEARCH INPUT LOCK & RE-TYPING BAN -> CRITICAL
- YOU ARE STRICTLY FORBIDDEN from using 'type' or 'type_and_enter' on the header search box if you are ALREADY on a search results page (or if the URL contains '/search/').
- DO NOT re-type the search keyword into the search bar if search results are already being displayed on screen. Re-typing into the search bar resets the page navigation and causes loops.
- If you need to see more search results, use 'scroll_down' or click filter tabs like 'People'. Do NOT click or type into the search bar again!

RULE 10: 2D VISUAL ASCII MAP READING & OVERLAY PRIORITY -> CRITICAL
- Read the 2D VISUAL ASCII LAYER MAP carefully.
- If a '🚨 FOREGROUND ACTIVE POPUP MODAL OVERLAY' section is present at the top of the map, prioritize interacting with the active modal buttons ([1], [2], etc.) to submit or dismiss the popup.
- Do NOT click elements flagged as '⚠️ COVERED UNDER POPUP' while an active modal popup is open in focus.

[CONSTRAINTS & REQUIRED KEYS] -> EXTREMELY IMPORTANT
- 'click', 'type', 'type_and_enter', 'hover', and 'upload_file' MUST have an integer "element_id".
- 'click_coordinates' MUST have integer "x" and "y" keys representing coordinates on the 1280x800 screen.
- STRICT NUMERIC IDs ONLY: Element IDs are ALWAYS integers (e.g., 1, 4, 15). NEVER output strings like "MISSING" or "search". If you can't find the ID in PAGE STATE, use 'wait' or 'scroll_down'.
- 'upload_file' MUST include a "file_path" key with the EXACT path provided in the GOAL. 
- 'type' and 'type_and_enter' MUST include a "text" key with the EXACT text from the GOAL.

[OUTPUT_FORMAT]
OUTPUT EXACTLY AND ONLY RAW VALID JSON. DO NOT WRAP IN MARKDOWN. NO COMMENTS. NO TRAILING COMMAS. SINGLE DICTIONARY ONLY. ALWAYS include "wait_time" except when action is 'task_done'.
Example 1: {{"thought": "Typing the description.", "action": "type", "element_id": 5, "text": "Product review", "wait_time": 2}}
Example 2: {{"thought": "Clicking publish button and waiting for submission.", "action": "click", "element_id": 12, "wait_time": 4}}
Example 3: {{"thought": "The page shows a success message. My goal is complete.", "action": "task_done"}}
Example 4: {{"thought": "A tutorial popup is blocking the view. Clicking the 'X' button at top-right of the popup at coordinates (900, 200) to dismiss it.", "action": "click_coordinates", "x": 900, "y": 200, "wait_time": 3}}"""
        },
        {
            "role": "user",
            "content": (
                f"GOAL: {goal}\n"
                f"{history_str}\n"
                f"{progress_str}\n"
                f"{page_snapshot}\n\n"
                f"Analyze the 2D Visual ASCII Layer Map above carefully. Check foreground active modal overlay buttons, uncovered elements, and covered elements before selecting your next action.\n"
                f"Action {attempt}/{max_attempts}. Give ONE action JSON."
            )
        }
    ]

    action_data = None
    for attempt_idx in range(1, 4):
        try:
            print(f"🧠 [DECIDE ACTION] LLM Call Attempt {attempt_idx}/3...")
            raw = ai_interpret.get_gemma_response_with_fallback(messages, max_new_tokens=1000, temperature=0.1, image_path=image_path, force_lightning_ai=force_lightning_ai)
            action_data = _extract_json(raw)
            if not _is_failed_action(action_data):
                break
            else:
                print(f"⚠️ [DECIDE ACTION] Failed to parse JSON or invalid action on attempt {attempt_idx}. Retrying...")
        except Exception as e:
            print(f"⚠️ [DECIDE ACTION] Exception on attempt {attempt_idx}: {e}")
            action_data = None
            
    if not action_data or _is_failed_action(action_data):
        action_data = {"action": "scroll_down", "wait_time": 2, "thought": "Fallback action due to LLM failure to return valid JSON."}

    # === UNIFIED PATH CORRECTION SAFETY NET ===
    if action_data and action_data.get("action") == "upload_file":
        # AI might hallucinate or truncate paths. Let's correct it using the verified path from the user goal/prompt.
        import re
        goal_paths = re.findall(r'([a-zA-Z]:\\[^\'\"\n\r]+?\.(?:png|jpg|jpeg|webp|gif|txt|csv|pdf))', goal, re.IGNORECASE)
        goal_paths += re.findall(r'([a-zA-Z]:/[^\'\"\n\r]+?\.(?:png|jpg|jpeg|webp|gif|txt|csv|pdf))', goal, re.IGNORECASE)
        goal_paths = [p.strip("'\" \t") for p in goal_paths]
        
        correct_path = None
        for gp in goal_paths:
            gp_clean = re.sub(r'\\+', r'\\', gp)
            if os.path.exists(gp_clean):
                correct_path = gp_clean
                break
                
        if correct_path:
            ai_path = action_data.get("file_path", "") or action_data.get("path", "")
            ai_path_clean = re.sub(r'\\+', r'\\', ai_path).strip("'\" \t") if ai_path else ""
            
            # If the AI-provided path doesn't exist or is different/truncated, override it
            if not ai_path_clean or not os.path.exists(ai_path_clean) or ai_path_clean.lower() != correct_path.lower():
                print(f"🩹 [PATH SAFETY NET] Correcting AI path '{ai_path}' to exact existing path from goal: '{correct_path}'")
                action_data["file_path"] = correct_path
                if "path" in action_data:
                    action_data["path"] = correct_path
                    
    return action_data

def verify_form_with_gemma(goal: str, page_snapshot: str) -> list:
    """
    Checks if the required fields are properly filled before publishing.
    Returns a list of missing or empty field names.
    """
    # Detect site
    site = "pinterest"
    if "facebook" in goal.lower() or "facebook.com" in page_snapshot.lower():
        site = "facebook"
        
    if site == "facebook":
        form_instructions = """Check if the Facebook Post Composer form is properly filled.
Look for the Post Description/Text and Image upload status in the PAGE STATE.
Fields to check: "description" (the main post text) and "image_upload"."""
        expected_json = """{
  "is_ready_to_publish": true/false,
  "missing_fields": ["description", "image_upload"]
}"""
    else:
        form_instructions = """Check if the Pinterest Pin creation form is properly filled according to the GOAL.
Look for the Title, Description, Link, and Image upload status in the PAGE STATE."""
        expected_json = """{
  "is_ready_to_publish": true/false,
  "missing_fields": ["title", "description", "link", "image_upload"]
}"""

    messages = [
        {
            "role": "user",
            "content": f"""You are a strict QA tester for an automated web agent.
GOAL: {goal}
PAGE STATE:
{page_snapshot}

{form_instructions}
If a field says [EMPTY], or if it is entirely missing from the PAGE STATE, it is NOT filled.
If it says [FILE ALREADY UPLOADED] or [ALREADY FILLED: "..."], it IS filled.

Output EXACTLY AND ONLY STRICT JSON with your findings:
{expected_json}
Do NOT write any introduction or explanation. Do NOT write markdown blocks. Just output raw JSON."""
        }
    ]
    
    fallback_fields = ["description", "image_upload"] if site == "facebook" else ["title", "description", "link", "image_upload"]
    
    data = {}
    for attempt_idx in range(1, 4):
        try:
            print(f"🧪 [VERIFY FORM] LLM Call Attempt {attempt_idx}/3...")
            raw = ai_interpret.get_gemma_response_with_fallback(messages, max_new_tokens=500, temperature=0.1)
            data = _extract_clean_json(raw)
            if data and ("is_ready_to_publish" in data or "missing_fields" in data):
                break
            else:
                print(f"⚠️ [VERIFY FORM] Invalid response format on attempt {attempt_idx}. Retrying...")
        except Exception as e:
            print(f"⚠️ [VERIFY FORM] Exception on attempt {attempt_idx}: {e}")
            data = {}
            
    if not data:
        return fallback_fields
        
    if not data.get("is_ready_to_publish", False):
        return data.get("missing_fields", fallback_fields)
    return []

def get_element_mapping(goal: str, page_snapshot: str, expected_fields: list) -> dict:
    """
    Asks Gemma to look at the first page snapshot and map required fields to their ai-ids.
    Uses Chain of Thought (CoT) by forcing Gemma to output reasoning and tag info before the ID.
    """
    fields_str = ",\n  ".join([f'"{f}": <object_or_null>' for f in expected_fields])
    clues_str = """- 'image_upload' is usually an <input type="file"> or an element mentioning "upload" or "photos or videos".
- 'title' is usually an input/textarea with a hint like "Add a title".
- 'description' is usually a textarea or role="textbox" with a hint like "Tell everyone" or "What's on your mind".
- 'link' is usually an input with a hint like "link" or "destination".
- 'publish_button' is usually a button with text "Publish", "Save", "Post", or "Share" (the header publish button).
- 'publish_now_button' is the final button on the preview/submission page with text like "Publish now", "Publish", or "Submit"."""

    messages = [
        {
            "role": "user",
            "content": f"""You are an expert UI element locator.
GOAL: {goal}
PAGE STATE:
{page_snapshot}

Analyze the PAGE STATE and map the required fields to their exact numeric 'ai-id'.
To ensure 99% accuracy, DO NOT just guess the ID. For each required field, you MUST provide:
1. "dom_line": The exact text of the line from the PAGE STATE.
2. "reasoning": Why this line matches the field.
3. "id": The integer ai-id.

Clues to help you find them:
{clues_str}

If an element is clearly visible in the PAGE STATE, provide its object. 
If an element is NOT visible on the page right now, output null.

OUTPUT EXACTLY AND ONLY STRICT JSON.
Example Output format:
{{
  "image_upload": {{
     "dom_line": "[15] <input file>: [EMPTY FILE INPUT] (Hint: \"Drag and drop\")",
     "reasoning": "This is a file input meant for uploading the image.",
     "id": 15
  }},
  "title": {{
     "dom_line": "[18] <textarea >: [EMPTY] (Hint: \"Add a title\")",
     "reasoning": "The hint explicitly says 'Add a title'.",
     "id": 18
  }},
  "description": null,
  "publish_button": {{
     "dom_line": "[30] <button >: Publish",
     "reasoning": "This is the primary publish button.",
     "id": 30
  }}
}}

Now output the exact JSON for the current PAGE STATE:
{{
  {fields_str}
}}
Do NOT write any introduction or explanation. Do NOT write markdown blocks. Just output raw JSON."""
        }
    ]
    
    mapping = {}
    for attempt_idx in range(1, 4):
        try:
            print(f"🔍 [ELEMENT MAPPING] LLM Call Attempt {attempt_idx}/3...")
            raw = ai_interpret.get_gemma_response_with_fallback(messages, max_new_tokens=1000, temperature=0.1)
            mapping = _extract_clean_json(raw)
            if mapping:
                break
            else:
                print(f"⚠️ [ELEMENT MAPPING] Invalid mapping response on attempt {attempt_idx}. Retrying...")
        except Exception as e:
            print(f"⚠️ [ELEMENT MAPPING] Exception on attempt {attempt_idx}: {e}")
            mapping = {}
            
    return mapping