autoair 0.1.0__py3-none-any.whl

autoair/__init__.py

@@ -0,0 +1,3 @@
+"""AutoAir — type-safe, extensible no-code DAG builder for Apache Airflow 3."""
+
+__version__ = "0.1.0"

autoair/_compat.py

@@ -0,0 +1,96 @@
+"""Airflow version detection and import shims (Airflow >= 3.0)."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any
+
+
+@lru_cache(maxsize=1)
+def airflow_version() -> tuple[int, int, int] | None:
+    """Installed Airflow version, or ``None`` when Airflow is not importable."""
+    try:
+        from importlib.metadata import version
+
+        raw = version("apache-airflow-core")
+    except Exception:
+        try:
+            from importlib.metadata import version
+
+            raw = version("apache-airflow")
+        except Exception:
+            return None
+    parts: list[int] = []
+    for piece in raw.split(".")[:3]:
+        digits = "".join(ch for ch in piece if ch.isdigit())
+        parts.append(int(digits or 0))
+    while len(parts) < 3:
+        parts.append(0)
+    return parts[0], parts[1], parts[2]
+
+
+def airflow_version_str() -> str | None:
+    v = airflow_version()
+    return ".".join(map(str, v)) if v else None
+
+
+def supports_react_apps() -> bool:
+    """``react_apps`` / ``external_views`` plugin attributes were added in Airflow 3.1."""
+    v = airflow_version()
+    return v is not None and v >= (3, 1, 0)
+
+
+def dagbag_class() -> Any:
+    """``DagBag`` moved to ``airflow.dag_processing.dagbag`` in Airflow 3.2."""
+    try:
+        from airflow.dag_processing.dagbag import DagBag
+    except ImportError:
+        from airflow.models.dagbag import DagBag
+    return DagBag
+
+
+def base_operator_classes() -> tuple[type, ...]:
+    """Operator base classes (in Airflow 3 every operator derives from the Task SDK one)."""
+    try:
+        from airflow.sdk import BaseOperator
+    except ImportError:
+        return ()
+    return (BaseOperator,)
+
+
+# Where Task SDK symbols live per Airflow version. Timetables moved into ``airflow.sdk``
+# in 3.2; deadline classes appeared in 3.1 under ``airflow.sdk.definitions.deadline``.
+# Entries: symbol -> [(min_version, module)], newest first.
+_SYMBOL_LOCATIONS: dict[str, list[tuple[tuple[int, int, int], str]]] = {
+    **{
+        name: [((3, 2, 0), "airflow.sdk"), ((3, 0, 0), module)]
+        for name, module in {
+            "CronTriggerTimetable": "airflow.timetables.trigger",
+            "DeltaTriggerTimetable": "airflow.timetables.trigger",
+            "MultipleCronTriggerTimetable": "airflow.timetables.trigger",
+            "CronDataIntervalTimetable": "airflow.timetables.interval",
+            "DeltaDataIntervalTimetable": "airflow.timetables.interval",
+            "EventsTimetable": "airflow.timetables.events",
+            "AssetOrTimeSchedule": "airflow.timetables.assets",
+        }.items()
+    },
+    **{
+        name: [((3, 2, 0), "airflow.sdk"), ((3, 1, 0), "airflow.sdk.definitions.deadline")]
+        for name in ("DeadlineAlert", "DeadlineReference", "AsyncCallback")
+    },
+    "AssetTriggeredTimetable": [
+        ((3, 4, 0), "airflow.sdk"),
+        ((3, 0, 0), "airflow.timetables.simple"),
+    ],
+}
+
+LATEST_KNOWN: tuple[int, int, int] = (3, 4, 0)
+
+
+def sdk_symbol_module(name: str, version: tuple[int, int, int] | None = None) -> str:
+    """Module to import an authoring symbol from for the given (or installed) Airflow."""
+    version = version or airflow_version() or LATEST_KNOWN
+    for min_version, module in _SYMBOL_LOCATIONS.get(name, []):
+        if version >= min_version:
+            return module
+    return "airflow.sdk"

autoair/_fs.py

@@ -0,0 +1,46 @@
+"""Small filesystem helpers shared by the state stores (templates, icons, settings)."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+
+def atomic_write_bytes(path: Path, data: bytes) -> None:
+    """Write ``path`` so readers see the old or the new content, never a torn file
+    (temp file in the same folder + ``os.replace``)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
+    try:
+        with os.fdopen(fd, "wb") as fh:
+            fh.write(data)
+        os.replace(tmp, path)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.unlink(tmp)
+        raise
+
+
+def atomic_write_text(path: Path, text: str) -> None:
+    atomic_write_bytes(path, text.encode("utf-8"))
+
+
+def _is_local(fs: Any) -> bool:
+    protocol = fs.protocol if isinstance(fs.protocol, (tuple, list)) else (fs.protocol,)
+    return "file" in protocol or "local" in protocol
+
+
+def fs_write_bytes(fs: Any, path: str, data: bytes) -> None:
+    """Write through an fsspec filesystem: atomically on local disks (object stores
+    replace objects atomically on their own)."""
+    if _is_local(fs):
+        atomic_write_bytes(Path(path), data)
+        return
+    fs.makedirs(path.rsplit("/", 1)[0], exist_ok=True)
+    fs.pipe_file(path, data)
+
+
+__all__ = ["atomic_write_bytes", "atomic_write_text", "fs_write_bytes"]

autoair/_subprocess.py

@@ -0,0 +1,50 @@
+"""Environment for child processes that execute user-authored DAG code.
+
+Import checks, Python DAG imports and test runs execute code written in the builder (or
+uploaded) on the AutoAir host. Those processes need Airflow's configuration, but never
+AutoAir's own credentials (git tokens, LLM API keys, the collaboration HMAC key...),
+which DAG authors must not be able to read from ``os.environ``.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Mapping
+from functools import lru_cache
+
+from autoair.sections import SandboxSettings, active
+
+
+@lru_cache(maxsize=8)
+def _secret_names(
+    patterns: tuple[str, ...], exact: frozenset[str]
+) -> tuple[tuple[re.Pattern[str], ...], frozenset[str]]:
+    return tuple(re.compile(p, re.I) for p in patterns if p), exact
+
+
+def is_secret_name(name: str) -> bool:
+    """``AUTOAIR_*`` settings holding secrets (nested ones included, e.g.
+    ``AUTOAIR_STORAGE__GIT__TOKEN``; ``AUTOAIR_SANDBOX__SECRET_ENV_PATTERN``) and provider
+    keys read by SDKs directly (``AUTOAIR_SANDBOX__SCRUB_ENV``)."""
+    cfg = active(SandboxSettings)
+    patterns, exact = _secret_names(
+        (cfg.secret_env_pattern, *cfg.secret_env_patterns_extra),
+        frozenset(n.upper() for n in (*cfg.scrub_env, *cfg.scrub_env_extra)),
+    )
+    upper = name.upper()
+    return upper in exact or (
+        upper.startswith("AUTOAIR_") and any(p.search(upper) for p in patterns)
+    )
+
+
+def user_code_env(
+    extra: Mapping[str, str] | None = None, base: Mapping[str, str] | None = None
+) -> dict[str, str]:
+    """``base`` (default ``os.environ``) without AutoAir secrets, plus ``extra``."""
+    env = {k: v for k, v in (os.environ if base is None else base).items() if not is_secret_name(k)}
+    env.update(extra or {})
+    return env
+
+
+__all__ = ["is_secret_name", "user_code_env"]

autoair/admin.py

@@ -0,0 +1,238 @@
+"""Admin panel back end: users seen, AI usage, pending proposals, cache and database status.
+
+The routes live in :mod:`autoair.api.routers.admin`; this module holds the logic so it is
+shared by the API and tests, and works on both state backends (files / SQL database).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import csv
+import datetime
+import hashlib
+import io
+import json
+import logging
+import threading
+import time
+from collections.abc import Iterable
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from autoair.audit import AuditRecord
+from autoair.persistence import document_store
+from autoair.persistence.documents import DocumentStore
+
+log = logging.getLogger("autoair.admin")
+
+USER_SEEN_INTERVAL = 300.0  # seconds between two "seen" writes for one user and process
+USAGE_SCAN_LIMIT = 20000  # AI audit records aggregated on the file backend
+
+
+class SeenUser(BaseModel):
+    user: str
+    roles: list[str] = Field(default_factory=list, description="AutoAir roles when last seen.")
+    permissions: int = Field(default=0, description="Number of AutoAir permissions held.")
+    auth: str | None = None
+    last_seen: str | None = None
+    actions: int = Field(default=0, description="Audit records (recent window).")
+    last_action: str | None = None
+    last_action_at: str | None = None
+
+
+class UsageRow(BaseModel):
+    user: str | None
+    provider: str | None = None
+    model: str | None = None
+    requests: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost_usd: float | None = None
+    errors: int = 0
+
+
+# --------------------------------------------------------------------------- users seen
+
+
+def users_store(settings: Any) -> DocumentStore:
+    return document_store(settings, "users_seen", file_url=str(settings.state_dir / "users"))
+
+
+def _user_key(user: str) -> str:
+    return hashlib.sha256(user.encode()).hexdigest()[:40]
+
+
+class SeenUsers:
+    """Remembers who used AutoAir and with which roles (throttled per user and process)."""
+
+    def __init__(self, store: DocumentStore, interval: float = USER_SEEN_INTERVAL) -> None:
+        self.store = store
+        self.interval = interval
+        self._last: dict[str, float] = {}
+        self._lock = threading.Lock()
+
+    def due(self, user: str, now: float | None = None) -> bool:
+        now = time.monotonic() if now is None else now
+        with self._lock:
+            if now - self._last.get(user, -1e18) < self.interval:
+                return False
+            self._last[user] = now
+            if len(self._last) > 10_000:
+                self._last.clear()
+            return True
+
+    def record(self, user: str, *, roles: Iterable[str], permissions: int, auth: str) -> None:
+        payload = {
+            "user": user,
+            "roles": sorted(set(roles)),
+            "permissions": permissions,
+            "auth": auth,
+            "last_seen": datetime.datetime.now(datetime.timezone.utc).isoformat(),
+        }
+        self.store.put(_user_key(user), json.dumps(payload).encode())
+
+    def all(self) -> list[SeenUser]:
+        out: list[SeenUser] = []
+        for key in self.store.keys():  # noqa: SIM118 - a DocumentStore, not a dict
+            found = self.store.get(key)
+            if found is None:
+                continue
+            try:
+                out.append(SeenUser.model_validate_json(found.data))
+            except ValueError:
+                continue
+        return out
+
+
+def seen_users_for(services: Any) -> SeenUsers:
+    existing = getattr(services, "_autoair_seen_users", None)
+    if isinstance(existing, SeenUsers):
+        return existing
+    created = SeenUsers(users_store(services.settings))
+    with contextlib.suppress(AttributeError, TypeError):
+        services._autoair_seen_users = created
+    return created
+
+
+def users_overview(seen: list[SeenUser], audit: list[AuditRecord]) -> list[SeenUser]:
+    """Users seen by the API merged with the users found in ``audit`` (newest first)."""
+    by_user: dict[str, SeenUser] = {u.user: u.model_copy() for u in seen}
+    for rec in audit:
+        if not rec.user:
+            continue
+        item = by_user.setdefault(rec.user, SeenUser(user=rec.user))
+        item.actions += 1
+        if item.last_action_at is None or rec.timestamp > item.last_action_at:
+            item.last_action_at = rec.timestamp
+            item.last_action = rec.action
+    return sorted(
+        by_user.values(),
+        key=lambda u: max(u.last_seen or "", u.last_action_at or ""),
+        reverse=True,
+    )
+
+
+# --------------------------------------------------------------------------- audit export
+
+CSV_FIELDS = ("timestamp", "action", "user", "dag_id", "details")
+
+
+def audit_csv(records: Iterable[AuditRecord]) -> str:
+    buf = io.StringIO()
+    writer = csv.writer(buf)
+    writer.writerow(CSV_FIELDS)
+    for rec in records:
+        writer.writerow(
+            [
+                rec.timestamp,
+                _cell(rec.action),
+                _cell(rec.user or ""),
+                _cell(rec.dag_id or ""),
+                _cell(json.dumps(rec.details, sort_keys=True, default=str)),
+            ]
+        )
+    return buf.getvalue()
+
+
+def _cell(text: str) -> str:
+    """Neutralise spreadsheet formulas (CSV injection): a leading = + - @ becomes text."""
+    return f"'{text}" if text[:1] in ("=", "+", "-", "@", "\t", "\r") else text
+
+
+# --------------------------------------------------------------------------- AI usage
+
+
+def aggregate_usage(records: Iterable[Any]) -> list[UsageRow]:
+    """Requests, tokens and cost per (user, provider, model) of AI audit records."""
+    rows: dict[tuple[str | None, str | None, str | None], UsageRow] = {}
+    for rec in records:
+        key = (rec.user, rec.provider, rec.model)
+        row = rows.setdefault(key, UsageRow(user=rec.user, provider=rec.provider, model=rec.model))
+        usage = rec.usage
+        row.requests += 1
+        row.input_tokens += usage.input_tokens + usage.cache_read_tokens + usage.cache_write_tokens
+        row.output_tokens += usage.output_tokens
+        if usage.cost_usd is not None:
+            row.cost_usd = round((row.cost_usd or 0.0) + usage.cost_usd, 8)
+        if rec.outcome != "ok":
+            row.errors += 1
+    return sorted(rows.values(), key=lambda r: (-(r.cost_usd or 0), -r.requests))
+
+
+def usage_rows(raw: Iterable[dict[str, Any]]) -> list[UsageRow]:
+    out = [
+        UsageRow(
+            user=r.get("user"),
+            provider=r.get("provider"),
+            model=r.get("model"),
+            requests=int(r.get("requests") or 0),
+            input_tokens=int(r.get("input_tokens") or 0),
+            output_tokens=int(r.get("output_tokens") or 0),
+            cost_usd=float(r["cost_usd"]) if r.get("cost_usd") is not None else None,
+            errors=int(r.get("errors") or 0),
+        )
+        for r in raw
+    ]
+    return sorted(out, key=lambda r: (-(r.cost_usd or 0), -r.requests))
+
+
+# --------------------------------------------------------------------------- proposals
+
+
+def all_pending_proposals(store: Any) -> list[Any]:
+    """Pending proposals of every DAG (SQL store: one query; file store: a scan)."""
+    if store is None:
+        return []
+    listing = getattr(store, "all_pending", None)
+    if callable(listing):
+        return list(listing())
+    fs = getattr(store, "fs", None)
+    root = getattr(store, "root", None)
+    if fs is None or root is None:
+        return []
+    try:
+        dirs = fs.ls(root, detail=False)
+    except (FileNotFoundError, OSError):
+        return []
+    out: list[Any] = []
+    for path in dirs:
+        dag_id = str(path).rstrip("/").rsplit("/", 1)[-1]
+        try:
+            out.extend(store.pending(dag_id))
+        except Exception:
+            log.debug("skipping proposals of %s", dag_id, exc_info=True)
+    return sorted(out, key=lambda p: p.timestamp, reverse=True)
+
+
+__all__ = [
+    "SeenUser",
+    "SeenUsers",
+    "UsageRow",
+    "aggregate_usage",
+    "all_pending_proposals",
+    "audit_csv",
+    "seen_users_for",
+    "usage_rows",
+    "users_overview",
+]

autoair/ai/__init__.py

@@ -0,0 +1,99 @@
+"""AutoAir AI: DAG generation, review / leak scan / optimise / fix features, retrieval
+(RAG), research, writers, guardrails, memory and tool-using agents on pydantic-ai.
+
+Everything is imported lazily (PEP 562): ``autoair.settings`` imports
+:mod:`autoair.ai.config` at start-up, and the plugin / server must start without the
+optional ``autoair[ai]`` extra (pydantic-ai).
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from autoair.ai.findings import Finding, FindingsReport, PatchOp, apply_findings, apply_patch
+    from autoair.ai.generator import (
+        GenerationResult,
+        describe_dag,
+        explain_problems,
+        generate_spec,
+        rank_operators,
+        select_operators,
+    )
+    from autoair.ai.llm import (
+        DEFAULT_MODEL,
+        LLM,
+        ChatMessage,
+        LLMError,
+        ModelConfig,
+        build_embedder,
+        build_llm,
+        provider_config,
+    )
+    from autoair.ai.problems import SpecCheck, SpecProblem, check_spec
+    from autoair.ai.usage import Usage, track_usage
+
+_EXPORTS: dict[str, str] = {
+    "Finding": "autoair.ai.findings",
+    "FindingsReport": "autoair.ai.findings",
+    "PatchOp": "autoair.ai.findings",
+    "apply_findings": "autoair.ai.findings",
+    "apply_patch": "autoair.ai.findings",
+    "GenerationResult": "autoair.ai.generator",
+    "describe_dag": "autoair.ai.generator",
+    "explain_problems": "autoair.ai.generator",
+    "generate_spec": "autoair.ai.generator",
+    "rank_operators": "autoair.ai.generator",
+    "select_operators": "autoair.ai.generator",
+    "DEFAULT_MODEL": "autoair.ai.llm",
+    "LLM": "autoair.ai.llm",
+    "ChatMessage": "autoair.ai.llm",
+    "LLMError": "autoair.ai.llm",
+    "ModelConfig": "autoair.ai.llm",
+    "build_embedder": "autoair.ai.llm",
+    "build_llm": "autoair.ai.llm",
+    "provider_config": "autoair.ai.llm",
+    "SpecCheck": "autoair.ai.problems",
+    "SpecProblem": "autoair.ai.problems",
+    "check_spec": "autoair.ai.problems",
+    "Usage": "autoair.ai.usage",
+    "track_usage": "autoair.ai.usage",
+}
+
+
+def __getattr__(name: str) -> Any:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module 'autoair.ai' has no attribute {name!r}")
+    value = getattr(importlib.import_module(module), name)
+    globals()[name] = value
+    return value
+
+
+__all__ = [
+    "DEFAULT_MODEL",
+    "LLM",
+    "ChatMessage",
+    "Finding",
+    "FindingsReport",
+    "GenerationResult",
+    "LLMError",
+    "ModelConfig",
+    "PatchOp",
+    "SpecCheck",
+    "SpecProblem",
+    "Usage",
+    "apply_findings",
+    "apply_patch",
+    "build_embedder",
+    "build_llm",
+    "check_spec",
+    "describe_dag",
+    "explain_problems",
+    "generate_spec",
+    "provider_config",
+    "rank_operators",
+    "select_operators",
+    "track_usage",
+]