main_back.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
import logging
import asyncio
from typing import List
import google.generativeai as genai
from dotenv import load_dotenv
import os

load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Sequential Test Step Generator API")

# Add CORS middleware to allow frontend requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class ModelInference:
    def __init__(self):
        """Initialize the model."""
        genai.configure(api_key=os.getenv("KEY"))
        self.model = genai.GenerativeModel("gemini-2.5-flash")
        self.device = "cuda"
        logger.info("Model loaded successfully!")

    def process_single_image(self, image: Image.Image) -> Image.Image:
        """Convert image to RGB PIL Image."""
        if image.mode != "RGB":
            image = image.convert("RGB")
        return image

    def predict_next_step_with_history(
        self, image: Image.Image, goal: str, completed_steps: List[str] = None
    ) -> str:
        """Predict the next step."""
        try:
            if completed_steps is None:
                completed_steps = []

            image = self.process_single_image(image)

            if completed_steps:
                history_str = "\n".join(
                    [f"{i + 1}. {step}" for i, step in enumerate(completed_steps)]
                )
                prompt = f"""Analyze this UI and generate the next test step.

Task: {goal}

Completed:
{history_str}

Output format: "ACTION: description [x1, y1, x2, y2]"
Actions: CLICK, TYPE, SCROLL, WAIT, VERIFY, SELECT, DRAG
Coordinates: normalized 0.0-1.0

Next step only:"""
            else:
                prompt = f"""Analyze this UI and generate the first test step.

Task: {goal}

Output format: "ACTION: description [x1, y1, x2, y2]"
Actions: CLICK, TYPE, SCROLL, WAIT, VERIFY, SELECT, DRAG
Coordinates: normalized 0.0-1.0

First step only:"""

            logger.info(
                f"Generating prediction with {len(completed_steps)} history steps"
            )
            response = self.model.generate_content([prompt, image])
            prediction = response.text.strip()
            logger.info(f"Generated prediction: {prediction}")
            return prediction

        except Exception as e:
            logger.error(f"Error during prediction: {str(e)}")
            raise

    def generate_step_sequence(
        self,
        image: Image.Image,
        task_description: str,
        action_history: str = "",
        max_steps: int = 10,
    ) -> List[str]:
        """Generate sequence of steps."""
        logger.info("Using recursive history-aware workflow generation")
        return self.generate_recursive_workflow(
            image=image,
            goal=task_description,
            initial_history=action_history,
            max_steps=max_steps,
        )

    def generate_recursive_workflow(
        self,
        image: Image.Image,
        goal: str,
        initial_history: str = "",
        max_steps: int = 10,
    ) -> List[str]:
        """Generate all workflow steps at once (faster)."""
        completed_steps = []

        if initial_history and initial_history.strip():
            if "→" in initial_history:
                completed_steps = [
                    s.strip() for s in initial_history.split("→") if s.strip()
                ]
            elif "," in initial_history:
                completed_steps = [
                    s.strip() for s in initial_history.split(",") if s.strip()
                ]
            else:
                completed_steps = [initial_history.strip()]

        logger.info(f"Generating all workflow steps at once for goal: {goal}")
        logger.info(f"Initial history: {completed_steps}")

        # Generate all steps in one call
        image = self.process_single_image(image)

        if completed_steps:
            history_str = "\n".join(
                [f"{i + 1}. {step}" for i, step in enumerate(completed_steps)]
            )
            prompt = f"""Analyze this UI and generate ALL remaining test steps to complete the task.

Task: {goal}

Already completed steps:
{history_str}

Generate the REMAINING steps needed to complete the task.

CRITICAL RULES:
- Output ONLY the steps, NO explanations, NO reasoning, NO extra text
- One step per line
- Format: "ACTION: description [x1, y1, x2, y2]"
- Actions: CLICK, TYPE, SCROLL, WAIT, VERIFY, SELECT, DRAG
- Coordinates: normalized 0.0-1.0
- For TYPE actions, describe what to type WITHOUT providing example values (e.g., "TYPE: Enter username in email field" NOT "TYPE: test@example.com")
- For CLICK actions, describe what to click (e.g., "CLICK: Click on the username input field")
- Maximum {max_steps} steps

Steps:"""
        else:
            prompt = f"""Analyze this UI and generate ALL test steps to complete the task.

Task: {goal}

Generate a complete sequence of steps to accomplish this task.

CRITICAL RULES:
- Output ONLY the steps, NO explanations, NO reasoning, NO extra text
- One step per line
- Format: "ACTION: description [x1, y1, x2, y2]"
- Actions: CLICK, TYPE, SCROLL, WAIT, VERIFY, SELECT, DRAG
- Coordinates: normalized 0.0-1.0
- For TYPE actions, describe what to type WITHOUT providing example values (e.g., "TYPE: Enter username in email field" NOT "TYPE: test@example.com")
- For CLICK actions, describe what to click (e.g., "CLICK: Click on the username input field")
- Maximum {max_steps} steps

Steps:"""

        try:
            logger.info("Generating all steps in single API call...")
            response = self.model.generate_content([prompt, image])
            all_steps_text = response.text.strip()

            # Parse steps (split by newlines)
            new_steps = []
            for line in all_steps_text.split("\n"):
                line = line.strip()
                # Skip empty lines, numbered prefixes, and explanatory text
                if not line:
                    continue
                # Remove numbering if present (e.g., "1. " or "1) ")
                if line and line[0].isdigit():
                    line = line.split(".", 1)[-1].strip()
                    line = line.split(")", 1)[-1].strip()
                # Only keep lines that start with action keywords
                if any(
                    line.upper().startswith(action)
                    for action in [
                        "CLICK:",
                        "TYPE:",
                        "SCROLL:",
                        "WAIT:",
                        "VERIFY:",
                        "SELECT:",
                        "DRAG:",
                    ]
                ):
                    new_steps.append(line)

            logger.info(f"Generated {len(new_steps)} steps in one call")
            for i, step in enumerate(new_steps):
                logger.info(f"Step {len(completed_steps) + i + 1}: {step}")

            return completed_steps + new_steps

        except Exception as e:
            logger.error(f"Error generating all steps: {str(e)}")
            raise


# Initialize model
logger.info("Initializing model inference...")
model_inference = ModelInference()
logger.info("Model inference ready!")


@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "running",
        "message": "Sequential Test Step Generator API",
        "device": str(model_inference.device),
        "model_loaded": True,
    }


@app.post("/predict")
async def predict(
    image: UploadFile = File(..., description="UI screenshot image"),
    action_history: str = Form(default="", description="Previous action history"),
    task_description: str = Form(..., description="Task description"),
    generate_sequence: bool = Form(
        default=True, description="Generate full sequence or single action"
    ),
):
    """Generate test steps based on UI image, action history, and task description."""
    try:
        await asyncio.sleep(0.5)

        image_data = await image.read()
        pil_image = Image.open(io.BytesIO(image_data))

        logger.info(f"Received image: {pil_image.size}, mode: {pil_image.mode}")
        logger.info(f"Task description: {task_description}")
        logger.info(
            f"Action history: {action_history[:100]}..."
            if action_history
            else "No history"
        )

        if generate_sequence:
            predicted_steps = model_inference.generate_step_sequence(
                image=pil_image,
                task_description=task_description,
                action_history=action_history,
                max_steps=10,
            )
        else:
            completed_steps = []
            if action_history and action_history.strip():
                if "→" in action_history:
                    completed_steps = [
                        s.strip() for s in action_history.split("→") if s.strip()
                    ]
                elif "," in action_history:
                    completed_steps = [
                        s.strip() for s in action_history.split(",") if s.strip()
                    ]
                else:
                    completed_steps = [action_history.strip()]

            predicted_action = model_inference.predict_next_step_with_history(
                image=pil_image, goal=task_description, completed_steps=completed_steps
            )
            predicted_steps = [predicted_action]

        return {
            "success": True,
            "steps": predicted_steps,
            "image_size": pil_image.size,
            "num_steps": len(predicted_steps),
        }

    except Exception as e:
        logger.error(f"Error processing request: {str(e)}", exc_info=True)
        return {"success": False, "error": "ERROR", "steps": []}


@app.get("/health")
async def health():
    """Detailed health check."""
    return {
        "status": "healthy",
        "device": str(model_inference.device),
        "model_loaded": model_inference.model is not None,
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)