selenium-python-ai-agent 0.2.0__py3-none-any.whl

selenium_agent/utils/config_manager.py

@@ -0,0 +1,81 @@
+"""
+CONFIG MANAGER
+==============
+Reads/writes .selenium-agent.json in the project root.
+Stores user preferences so they don't have to repeat CLI flags.
+
+Priority order (highest to lowest):
+  CLI flag > .env > .selenium-agent.json > built-in defaults
+"""
+
+import json
+import os
+from pathlib import Path
+from selenium_agent.utils.logger import setup_logger
+
+logger = setup_logger("Config")
+
+CONFIG_FILENAME = ".selenium-agent.json"
+
+DEFAULTS = {
+    "provider":  "anthropic",
+    "model":     None,
+    "headless":  False,
+    "mode":      "pytest",
+    "base_url":  None,   # persisted across runs — no need to pass --url repeatedly
+}
+
+
+def _find_config_path() -> Path:
+    """Look for config in CWD first, then home directory."""
+    cwd_config = Path.cwd() / CONFIG_FILENAME
+    if cwd_config.exists():
+        return cwd_config
+    return Path.home() / CONFIG_FILENAME
+
+
+def load() -> dict:
+    """Load config from file. Returns defaults if file doesn't exist."""
+    path = _find_config_path()
+    if not path.exists():
+        return dict(DEFAULTS)
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        merged = {**DEFAULTS, **data}
+        return merged
+    except Exception as e:
+        logger.warning(f"⚠️  Could not read {path}: {e} — using defaults")
+        return dict(DEFAULTS)
+
+
+def save(updates: dict, config_path: Path | None = None) -> Path:
+    """
+    Merge updates into existing config and save.
+    Creates config in CWD if it doesn't exist yet.
+    """
+    path = config_path or _find_config_path()
+    if not path.exists():
+        path = Path.cwd() / CONFIG_FILENAME
+
+    current = load()
+    current.update({k: v for k, v in updates.items() if v is not None})
+
+    path.write_text(json.dumps(current, indent=2), encoding="utf-8")
+    return path
+
+
+def get_effective(cli_args: dict) -> dict:
+    """
+    Merge config + CLI args. CLI always wins over config.
+    cli_args: dict of {key: value} from argparse (None means not provided).
+    """
+    config = load()
+    effective = dict(config)
+
+    for key, value in cli_args.items():
+        if value is not None and value is not False:
+            effective[key] = value
+        elif key == "headless" and value is True:
+            effective[key] = True
+
+    return effective

selenium_agent/utils/json_utils.py

@@ -0,0 +1,103 @@
+"""
+JSON UTILITIES
+==============
+Robust extraction of JSON objects from LLM responses.
+
+LLMs wrap JSON in markdown fences, add prose before/after, or truncate
+mid-object when they hit the token limit. Every agent (Planner, Coder,
+Healer) parses LLM JSON through this single module so the behavior is
+consistent and battle-tested in one place.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+
+class LLMJSONError(ValueError):
+    """Raised when no usable JSON object could be extracted from an LLM response."""
+
+    def __init__(self, message: str, raw: str = ""):
+        snippet = (raw or "")[:400]
+        super().__init__(f"{message}\n--- response snippet ---\n{snippet}")
+        self.raw = raw
+
+
+def strip_markdown_fences(raw: str) -> str:
+    """Remove ```json ... ``` fences, keeping the largest fenced block if present."""
+    if "```" not in raw:
+        return raw.strip()
+    blocks = re.findall(r"```(?:json)?\s*([\s\S]*?)```", raw)
+    if blocks:
+        # Take the largest block — small ones are usually inline examples
+        return max(blocks, key=len).strip()
+    # Unclosed fence: drop everything before the first fence marker
+    return raw.split("```", 1)[-1].removeprefix("json").strip()
+
+
+def find_balanced_object(raw: str) -> str:
+    """
+    Return the first balanced top-level {...} object in raw.
+    If the object is truncated, repair it structurally: close the open
+    string, then close arrays/objects in the reverse order they were opened.
+    """
+    start = raw.find("{")
+    if start == -1:
+        raise LLMJSONError("No JSON object found in LLM response", raw)
+
+    stack: list[str] = []
+    in_str, esc = False, False
+    for i, ch in enumerate(raw[start:], start):
+        if esc:
+            esc = False
+            continue
+        if ch == "\\" and in_str:
+            esc = True
+            continue
+        if ch == '"':
+            in_str = not in_str
+            continue
+        if in_str:
+            continue
+        if ch in "{[":
+            stack.append(ch)
+        elif ch in "}]":
+            if stack:
+                stack.pop()
+            if not stack:
+                return raw[start : i + 1]
+
+    # Truncated — repair by unwinding the open-structure stack
+    frag = raw[start:]
+    if in_str:
+        frag += '"'
+    for opener in reversed(stack):
+        frag += "}" if opener == "{" else "]"
+    return frag
+
+
+def extract_json_object(raw: str) -> dict:
+    """
+    Extract and parse a JSON object from an arbitrary LLM response.
+
+    Handles: markdown fences, leading/trailing prose, truncated output.
+    Raises LLMJSONError when nothing parseable can be recovered.
+    """
+    if not raw or not raw.strip():
+        raise LLMJSONError("Empty LLM response", raw)
+
+    candidate = find_balanced_object(strip_markdown_fences(raw))
+    try:
+        parsed = json.loads(candidate)
+    except json.JSONDecodeError:
+        # Last resort: remove trailing commas (a common LLM slip) and retry
+        cleaned = re.sub(r",\s*([}\]])", r"\1", candidate)
+        try:
+            parsed = json.loads(cleaned)
+        except json.JSONDecodeError as exc:
+            raise LLMJSONError(f"LLM returned invalid JSON: {exc}", raw) from exc
+
+    if not isinstance(parsed, dict):
+        raise LLMJSONError(f"Expected JSON object, got {type(parsed).__name__}", raw)
+    return parsed

selenium_agent/utils/llm.py

@@ -0,0 +1,240 @@
+"""
+LLM provider helpers for Anthropic and OpenAI.
+
+All agents talk to the LLM through this module, which provides:
+- provider normalization + API-key resolution
+- automatic retries with exponential backoff on transient failures
+  (rate limits, 5xx, timeouts, connection errors)
+- special handling for OpenAI reasoning models (gpt-5*, o*) whose
+  reasoning tokens count against max_output_tokens
+"""
+
+import os
+import time
+from importlib import import_module
+
+from selenium_agent.utils.logger import setup_logger
+
+logger = setup_logger("LLM")
+
+DEFAULT_PROVIDER = "anthropic"
+DEFAULT_MODELS = {
+    "anthropic": "claude-sonnet-5",
+    "openai": "gpt-5-mini",
+}
+API_KEY_ENV_VARS = {
+    "anthropic": "ANTHROPIC_API_KEY",
+    "openai": "OPENAI_API_KEY",
+}
+
+MAX_RETRIES = 3
+BACKOFF_BASE_SECONDS = 2.0
+
+# Error class-name fragments that must NOT be retried — the request itself is bad.
+_NON_RETRYABLE = (
+    "authentication", "permissiondenied", "notfound",
+    "badrequest", "invalidrequest", "unprocessable",
+)
+
+
+def normalize_provider(provider: str | None) -> str:
+    normalized = (provider or DEFAULT_PROVIDER).strip().lower()
+    if normalized not in API_KEY_ENV_VARS:
+        supported = ", ".join(sorted(API_KEY_ENV_VARS))
+        raise ValueError(f"Unsupported provider '{provider}'. Supported providers: {supported}")
+    return normalized
+
+
+def get_default_model(provider: str) -> str:
+    normalized = normalize_provider(provider)
+    return DEFAULT_MODELS[normalized]
+
+
+def get_api_key_env_var(provider: str) -> str:
+    normalized = normalize_provider(provider)
+    return API_KEY_ENV_VARS[normalized]
+
+
+def format_missing_api_key_error(provider: str) -> str:
+    normalized = normalize_provider(provider)
+    env_var = get_api_key_env_var(normalized)
+    provider_name = "Anthropic Claude" if normalized == "anthropic" else "OpenAI"
+    return (
+        f"{provider_name} API key required!\n"
+        f"Pass it as: SeleniumAgent(provider='{normalized}', api_key='your-key')\n"
+        f"Or set env var: export {env_var}='your-key'\n"
+        f"Get your key at: https://console.anthropic.com"
+    )
+
+
+def resolve_api_key(provider: str, api_key: str | None = None) -> str | None:
+    normalized = normalize_provider(provider)
+    if api_key:
+        return api_key
+    env_var = get_api_key_env_var(normalized)
+    return os.environ.get(env_var)
+
+
+def _is_retryable(exc: Exception) -> bool:
+    name = type(exc).__name__.lower()
+    if any(fragment in name for fragment in _NON_RETRYABLE):
+        return False
+    status = getattr(exc, "status_code", None)
+    if isinstance(status, int):
+        return status == 429 or status >= 500
+    # Rate limits, overload, timeouts, connection drops → retry
+    return any(f in name for f in ("ratelimit", "overloaded", "timeout", "connection",
+                                   "internalserver", "apierror", "apistatus"))
+
+
+class BaseLLMClient:
+    """Minimal interface shared by all provider implementations."""
+
+    model: str = ""
+
+    def generate_text(self, system_prompt: str, user_prompt: str, max_tokens: int,
+                      json_mode: bool = False) -> str:
+        """
+        Generate text with automatic retries on transient failures.
+
+        json_mode=True engages the provider's native structured-output
+        mechanism (OpenAI json_object format / Anthropic '{' prefill) so
+        responses are valid JSON instead of best-effort prose.
+        """
+        last_exc: Exception | None = None
+        for attempt in range(1, MAX_RETRIES + 1):
+            try:
+                return self._generate_once(system_prompt, user_prompt, max_tokens, json_mode)
+            except Exception as exc:  # noqa: BLE001 — provider SDKs raise many types
+                if not _is_retryable(exc) or attempt == MAX_RETRIES:
+                    raise
+                delay = BACKOFF_BASE_SECONDS * (2 ** (attempt - 1))
+                logger.warning(
+                    f"⚠️  LLM call failed ({type(exc).__name__}) — "
+                    f"retry {attempt}/{MAX_RETRIES - 1} in {delay:.0f}s"
+                )
+                time.sleep(delay)
+                last_exc = exc
+        raise last_exc  # pragma: no cover — unreachable
+
+    def _generate_once(self, system_prompt: str, user_prompt: str, max_tokens: int,
+                       json_mode: bool = False) -> str:
+        raise NotImplementedError
+
+
+class AnthropicLLMClient(BaseLLMClient):
+    """Anthropic Claude-backed text generation client."""
+
+    def __init__(self, api_key: str, model: str):
+        try:
+            anthropic = import_module("anthropic")
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                "The 'anthropic' package is required for provider='anthropic'. "
+                "Install it with: pip install anthropic"
+            ) from exc
+        self.client = anthropic.Anthropic(api_key=api_key)
+        self.model = model
+
+    def _generate_once(self, system_prompt: str, user_prompt: str, max_tokens: int,
+                       json_mode: bool = False) -> str:
+        messages = [{"role": "user", "content": user_prompt}]
+        if json_mode:
+            # Prefill trick: forcing the assistant to start at '{' makes
+            # Claude emit pure JSON with no prose or fences.
+            messages.append({"role": "assistant", "content": "{"})
+
+        message = self.client.messages.create(
+            model=self.model,
+            max_tokens=max_tokens,
+            system=system_prompt,
+            messages=messages,
+        )
+        parts = []
+        for content_block in message.content:
+            text = getattr(content_block, "text", None)
+            if text:
+                parts.append(text)
+        text = "\n".join(parts).strip()
+        if json_mode and not text.startswith("{"):
+            text = "{" + text
+        return text
+
+
+class OpenAILLMClient(BaseLLMClient):
+    """OpenAI-backed text generation client (Responses API)."""
+
+    def __init__(self, api_key: str, model: str):
+        try:
+            openai = import_module("openai")
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                "The 'openai' package is required for provider='openai'. "
+                "Install it with: pip install openai"
+            ) from exc
+        self.client = openai.OpenAI(api_key=api_key)
+        self.model = model
+
+    def _is_reasoning_model(self) -> bool:
+        m = self.model.lower()
+        return m.startswith(("gpt-5", "o1", "o3", "o4"))
+
+    def _generate_once(self, system_prompt: str, user_prompt: str, max_tokens: int,
+                       json_mode: bool = False) -> str:
+        params = {
+            "model": self.model,
+            "instructions": system_prompt,
+            "input": user_prompt,
+            "max_output_tokens": max_tokens,
+        }
+        if json_mode:
+            # Native structured output — the model cannot emit invalid JSON.
+            # The API requires the word 'json' in the input message itself.
+            params["text"] = {"format": {"type": "json_object"}}
+            if "json" not in user_prompt.lower():
+                params["input"] = user_prompt + "\n\nRespond with a valid JSON object."
+        if self._is_reasoning_model():
+            # Reasoning tokens count against max_output_tokens — keep effort low
+            # and leave headroom so the visible answer isn't starved.
+            params["reasoning"] = {"effort": "low"}
+            params["max_output_tokens"] = max(max_tokens * 2, 4000)
+
+        try:
+            response = self.client.responses.create(**params)
+        except TypeError:
+            # Older SDK without `reasoning`/`text` support
+            params.pop("reasoning", None)
+            params.pop("text", None)
+            response = self.client.responses.create(**params)
+
+        text = (response.output_text or "").strip()
+        if not text and getattr(response, "status", "") == "incomplete":
+            # Reasoning consumed the entire budget — one retry with a bigger cap
+            params["max_output_tokens"] = params["max_output_tokens"] * 2
+            response = self.client.responses.create(**params)
+            text = (response.output_text or "").strip()
+
+        if not text:
+            raise RuntimeError(
+                f"OpenAI model '{self.model}' returned an empty response "
+                f"(status={getattr(response, 'status', 'unknown')}). "
+                f"Try a larger max_tokens or a non-reasoning model like gpt-4o-mini."
+            )
+        return text
+
+
+def create_llm_client(
+    provider: str,
+    api_key: str,
+    model: str | None = None,
+) -> BaseLLMClient:
+    """Create a provider-specific LLM client."""
+    normalized = normalize_provider(provider)
+    resolved_model = model or get_default_model(normalized)
+
+    if normalized == "anthropic":
+        return AnthropicLLMClient(api_key=api_key, model=resolved_model)
+    if normalized == "openai":
+        return OpenAILLMClient(api_key=api_key, model=resolved_model)
+
+    raise ValueError(f"Unsupported provider '{provider}'")

selenium_agent/utils/logger.py

@@ -0,0 +1,37 @@
+"""
+Logger utility for Selenium AI Agent
+"""
+
+import logging
+import os
+from datetime import datetime
+
+
+def setup_logger(name: str, log_to_file: bool = False) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+
+    if logger.handlers:
+        return logger
+
+    formatter = logging.Formatter(
+        "%(asctime)s | %(name)s | %(levelname)s | %(message)s",
+        datefmt="%H:%M:%S"
+    )
+
+    # Console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    # Optional file handler
+    if log_to_file:
+        os.makedirs("logs", exist_ok=True)
+        log_file = f"logs/agent_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setLevel(logging.DEBUG)
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger

selenium_agent/utils/paths.py

@@ -0,0 +1,57 @@
+"""
+Path helpers for generated Selenium agent files.
+"""
+
+from pathlib import Path
+
+
+def get_output_root(output_dir: str) -> Path:
+    """Return the absolute root directory for generated files."""
+    return Path(output_dir).resolve()
+
+
+def resolve_input_path(file_path: str, output_dir: str) -> Path:
+    """
+    Resolve an input file path from either the current working directory or output_dir.
+
+    This supports both:
+    - generated_tests/tests/test_login.py
+    - tests/test_login.py when output_dir is generated_tests
+    """
+    candidate = Path(file_path)
+    if candidate.is_absolute():
+        return candidate.resolve()
+
+    output_root = get_output_root(output_dir)
+
+    # Paths that already include the output_dir prefix resolve against it
+    # first — avoids the CWD accidentally shadowing the intended target.
+    if candidate.parts and candidate.parts[0] == output_root.name:
+        prefixed = (output_root.parent / candidate).resolve()
+        if prefixed.exists():
+            return prefixed
+
+    if candidate.exists():
+        return candidate.resolve()
+
+    if candidate.parts and candidate.parts[0] == output_root.name:
+        return (output_root.parent / candidate).resolve()
+
+    return (output_root / candidate).resolve()
+
+
+def safe_output_path(output_dir: str, filename: str) -> Path:
+    """Resolve a model-provided filename and ensure it stays inside output_dir."""
+    relative_path = Path(filename)
+    if relative_path.is_absolute():
+        raise ValueError(f"Refusing to write absolute path from model output: {filename}")
+
+    output_root = get_output_root(output_dir)
+    destination = (output_root / relative_path).resolve()
+
+    try:
+        destination.relative_to(output_root)
+    except ValueError as exc:
+        raise ValueError(f"Refusing to write outside output_dir: {filename}") from exc
+
+    return destination

selenium_agent/utils/spec_writer.py

@@ -0,0 +1,126 @@
+"""
+SPEC WRITER
+===========
+Renders a structured test plan (the Planner's JSON) as a human-readable
+Markdown spec and persists both to the `specs/` directory — mirroring how
+Playwright's planner agent saves Markdown test plans.
+
+specs/
+├── login-page.md      ← human-readable plan (review / edit / re-run)
+└── login-page.json    ← machine-readable plan (input to the Generator)
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from datetime import datetime
+from pathlib import Path
+
+
+def slugify(text: str, max_len: int = 60) -> str:
+    """'Test the Login page!' → 'test-the-login-page'"""
+    slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
+    return slug[:max_len].rstrip("-") or "test-plan"
+
+
+def render_markdown(plan: dict, instruction: str = "") -> str:
+    """Render a plan dict (pytest or bdd mode) as a Markdown test plan."""
+    mode = plan.get("mode", "pytest")
+    url = plan.get("url", "")
+    scenarios = plan.get("test_scenarios") or plan.get("scenarios") or []
+    pages = plan.get("page_objects_needed", [])
+    locators = plan.get("locators", [])
+
+    lines = [
+        f"# Test Plan — {plan.get('feature_title') or instruction or 'Selenium Test Suite'}",
+        "",
+        f"- **Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
+        f"- **Mode**: {mode}",
+        f"- **Target URL**: {url or 'n/a'}",
+        f"- **Browser**: {plan.get('browser', 'chrome')} (headless={plan.get('headless', False)})",
+        "",
+        "## Page Objects",
+        "",
+    ]
+    for page in pages:
+        lines.append(f"- `{page}`")
+    if not pages:
+        lines.append("- _none planned_")
+
+    lines += ["", "## Scenarios", ""]
+    for idx, sc in enumerate(scenarios, 1):
+        name = sc.get("name", f"Scenario {idx}")
+        lines.append(f"### {sc.get('id', f'TC{idx:03d}')} — {name}")
+        lines.append("")
+        if sc.get("description"):
+            lines.append(f"{sc['description']}")
+            lines.append("")
+        lines.append("**Steps:**")
+        for step in sc.get("steps", []):
+            if isinstance(step, dict):  # BDD step: {"type": "Given", "text": "..."}
+                lines.append(f"1. **{step.get('type', '')}** {step.get('text', '')}")
+            else:
+                lines.append(f"1. {step}")
+        if sc.get("expected_result"):
+            lines.append("")
+            lines.append(f"**Expected:** {sc['expected_result']}")
+        if sc.get("test_data"):
+            lines.append("")
+            lines.append(f"**Test data:** `{json.dumps(sc['test_data'])}`")
+        lines.append("")
+
+    if locators:
+        lines += [
+            "## Locators (from live DOM scan)",
+            "",
+            "| Page Object | Element | Selector | Wait |",
+            "|---|---|---|---|",
+        ]
+        for loc in locators:
+            selector = loc.get("css") or loc.get("xpath") or ""
+            lines.append(
+                f"| {loc.get('page_object', '')} | {loc.get('element', '')} "
+                f"| `{selector}` | {loc.get('wait_condition', 'visible')} |"
+            )
+        lines.append("")
+
+    if plan.get("notes"):
+        lines += ["## Notes", "", str(plan["notes"]), ""]
+
+    return "\n".join(lines)
+
+
+def save_spec(plan: dict, instruction: str, specs_dir: str | Path = "specs") -> dict:
+    """
+    Save the plan as Markdown + JSON in specs_dir.
+    Returns {"markdown": path, "json": path}.
+    """
+    specs_root = Path(specs_dir)
+    specs_root.mkdir(parents=True, exist_ok=True)
+
+    slug = slugify(plan.get("feature_name") or instruction)
+    md_path = specs_root / f"{slug}.md"
+    json_path = specs_root / f"{slug}.json"
+
+    md_path.write_text(render_markdown(plan, instruction), encoding="utf-8")
+    json_path.write_text(json.dumps(plan, indent=2), encoding="utf-8")
+
+    return {"markdown": str(md_path), "json": str(json_path)}
+
+
+def load_plan(plan_file: str | Path) -> dict:
+    """Load a saved plan (.json) for the Generator to consume."""
+    path = Path(plan_file)
+    if not path.exists():
+        raise FileNotFoundError(f"Plan file not found: {path}")
+    if path.suffix == ".md":
+        # Allow pointing at the markdown twin — load the sibling .json
+        sibling = path.with_suffix(".json")
+        if not sibling.exists():
+            raise FileNotFoundError(
+                f"Markdown plans are for humans — the Generator needs the JSON twin, "
+                f"which was not found: {sibling}"
+            )
+        path = sibling
+    return json.loads(path.read_text(encoding="utf-8"))