SelfHealing.py

from transformers import LlamaForCausalLM, LlamaTokenizer
import torch
from playwright.sync_api import sync_playwright

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

# Define the input model for the request
class InputData(BaseModel):
    text: str
    
# Load LLaMA 2 model and tokenizer
model_name = "meta-llama/Llama-2-7b-hf"  # Use appropriate LLaMA 2 model variant
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaForCausalLM.from_pretrained(model_name)

# Example function to generate alternative locators using LLaMA 2
def generate_alternative_locators(html_content, failed_locator):
    prompt = f"""
    The following HTML structure contains a failing locator: "{failed_locator}".
    Suggest alternative XPaths or CSS selectors that could identify the same element.

    HTML:
    {html_content}
    """

    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=500)

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    suggestions = response.split("\n")  # Assuming each suggestion is on a new line

    return suggestions


# Example of self-healing process using Playwright and LLaMA 2 suggestions
@app.post("/predict")
def attempt_self_healing(page, failed_locator: str):
    # Retrieve HTML of the page
    html = page.content()

    # Get alternative locators from LLaMA 2
    suggestions = generate_alternative_locators(html, failed_locator)
    print("Suggested alternatives:", suggestions)

    # Attempt to click with alternative locators
    for locator in suggestions:
        try:
            page.locator(locator.strip()).click()
            print(f"Self-healing successful with locator: {locator}")
            return
        except Exception as e:
            print(f"Locator failed: {locator}, Error: {e}")

    print("Self-healing failed. No valid locator found.")