agent_common.py

import os
import pathlib
import time
import re
from dataclasses import dataclass
from typing import Optional

from dotenv import load_dotenv
from google import genai

ENV_PATH = os.path.join(os.path.dirname(__file__), ".env")
load_dotenv(dotenv_path=ENV_PATH, override=True)

DEFAULT_MODEL = os.getenv("GEMINI_MODEL", "models/gemini-2.5-flash")
API_KEY = os.getenv("GEMINI_API_KEY")

if not API_KEY:
    raise RuntimeError("Missing GEMINI_API_KEY. Put it in ~/evcore/agents/.env")

CLIENT = genai.Client(api_key=API_KEY)

SYSTEM_RULES = """You are an engineering agent working on the EV Core project on a Jetson AGX Orin.
Hard rules:
- Do not invent file contents. Use only the provided context.
- Be concise. Prefer bullet points and clear headings when writing non-code.
- Never include secrets (API keys, tokens).
- Follow the ROLE instructions exactly, especially output format rules.
"""

@dataclass
class RunPaths:
    run_dir: pathlib.Path
    context_path: pathlib.Path
    out_dir: pathlib.Path

def ensure_dir(p: pathlib.Path) -> None:
    p.mkdir(parents=True, exist_ok=True)

def read_text(p: pathlib.Path) -> str:
    return p.read_text(encoding="utf-8", errors="replace")

def write_text(p: pathlib.Path, s: str) -> None:
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(s, encoding="utf-8")

def model_name(override: Optional[str] = None) -> str:
    return override or DEFAULT_MODEL

def _parse_retry_delay_seconds(err_text: str) -> Optional[float]:
    """
    Try to extract retry delay from Gemini error text.
    Handles patterns like:
      - retryDelay': '27s'
      - Please retry in 7.245659367s
    """
    m = re.search(r"retryDelay'\s*:\s*'(\d+)s'", err_text)
    if m:
        return float(m.group(1))
    m = re.search(r"retry in ([0-9]+(?:\.[0-9]+)?)s", err_text, re.IGNORECASE)
    if m:
        return float(m.group(1))
    return None

def generate_text(prompt: str, model: Optional[str] = None) -> str:
    """
    Calls Gemini with lightweight retry for transient errors.
    Retries:
      - 503 UNAVAILABLE (high demand)
      - 429 RESOURCE_EXHAUSTED (rate limits / token limits)
    """
    chosen_model = model_name(model)

    last_exc: Optional[Exception] = None
    for attempt in range(1, 4):
        try:
            resp = CLIENT.models.generate_content(
                model=chosen_model,
                contents=[{"role": "user", "parts": [{"text": prompt}]}],
            )
            return (resp.text or "").strip()
        except Exception as e:
            last_exc = e
            msg = repr(e)

            transient = (
                ("503" in msg) or ("UNAVAILABLE" in msg) or
                ("429" in msg) or ("RESOURCE_EXHAUSTED" in msg)
            )
            if not transient:
                raise

            delay = _parse_retry_delay_seconds(msg)
            if delay is None:
                # small exponential-ish backoff
                delay = 2.0 * attempt

            time.sleep(delay)

    # If we got here, all retries failed
    raise last_exc if last_exc else RuntimeError("generate_text failed with unknown error")