dataforge-07 0.1.0__py3-none-any.whl

dataforge/env/observation.py

@@ -0,0 +1,61 @@
+"""Observation builder for the DataForge RL environment.
+
+Constructs agent-visible observations containing partial data views,
+scratchpad summaries, tool results, and step budget information.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+__all__ = ["DataForgeObservation", "ToolResult"]
+
+
+class ToolResult(BaseModel):
+    """Result of a single tool-use action.
+
+    Args:
+        action_type: The action type that produced this result.
+        success: Whether the action succeeded.
+        data: Action-specific result data (rows, stats, matches, etc.).
+        error: Structured error information if the action failed.
+    """
+
+    action_type: str
+    success: bool = True
+    data: Any = None
+    error: dict[str, Any] | None = None
+
+    model_config = {"frozen": True}
+
+
+class DataForgeObservation(BaseModel):
+    """Agent-visible observation returned after each environment step.
+
+    Args:
+        visible_rows: Dataset rows returned by INSPECT_ROWS or reset.
+        detector_hints: Optional hints from detectors (partial ground truth).
+        scratchpad_summary: Compact summary of the agent's scratchpad.
+        step_budget_remaining: Steps left before auto-finalize.
+        tool_usage_history: Last 5 tool results for context.
+        latest_result: Result of the most recent action.
+        done: Whether the episode has ended.
+        reward: Step reward.
+        cumulative_reward: Running total reward for the episode.
+        metadata: Additional key-value metadata.
+    """
+
+    visible_rows: list[dict[str, Any]] | None = None
+    detector_hints: list[str] | None = None
+    scratchpad_summary: str = ""
+    step_budget_remaining: int = 0
+    tool_usage_history: list[ToolResult] = Field(default_factory=list)
+    latest_result: ToolResult | None = None
+    done: bool = False
+    reward: float = 0.0
+    cumulative_reward: float = 0.0
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    model_config = {"frozen": True}

dataforge/env/openenv_core.py

@@ -0,0 +1,161 @@
+"""OpenEnv-core adapter for the DataForge RL environment."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from dataforge.env.environment import DataForgeEnv
+
+if TYPE_CHECKING:
+
+    class OpenEnvAction(BaseModel):
+        """Typed stand-in for openenv-core's action model."""
+
+        metadata: dict[str, Any] = Field(default_factory=dict)
+
+    class OpenEnvObservation(BaseModel):
+        """Typed stand-in for openenv-core's observation model."""
+
+        done: bool = False
+        reward: float | None = None
+        metadata: dict[str, Any] = Field(default_factory=dict)
+
+    class OpenEnvEnvironment:
+        """Typed stand-in for openenv-core's environment base."""
+
+        def __init__(self) -> None: ...
+
+    def create_app(*args: Any, **kwargs: Any) -> Any:
+        """Typed stand-in for openenv-core's FastAPI app factory."""
+        ...
+
+else:
+    try:
+        from openenv.core.env_server import Action as OpenEnvAction
+        from openenv.core.env_server import Environment as OpenEnvEnvironment
+        from openenv.core.env_server import Observation as OpenEnvObservation
+        from openenv.core.env_server import create_app
+    except ImportError as exc:  # pragma: no cover - exercised only without openenv extra
+        raise RuntimeError(
+            "The OpenEnv adapter requires the openenv extra: pip install 'dataforge[openenv]'."
+        ) from exc
+
+
+class DataForgeOpenEnvAction(OpenEnvAction):
+    """OpenEnv action wrapper for DataForge's typed action payloads."""
+
+    action_type: str = Field(min_length=1)
+    row_indices: list[int] | None = None
+    column_names: list[str] | None = None
+    query: str | None = None
+    sql: str | None = None
+    test_type: str | None = None
+    test: str | None = None
+    column: str | None = None
+    threshold: float | None = None
+    pattern: str | None = None
+    regex: str | None = None
+    expect_match: bool | None = None
+    claim: str | None = None
+    affected_rows: list[int] | None = None
+    affected_columns: list[str] | None = None
+    root_cause_type: str | None = None
+    error_indices: list[int] | None = None
+    row: int | None = None
+    issue_type: str | None = None
+    new_value: str | None = None
+    proposed_value: str | None = None
+    justification: str | None = None
+    fix_type: str | None = None
+
+    def as_dataforge_payload(self) -> dict[str, Any]:
+        """Return the action payload expected by ``DataForgeEnv.step``."""
+        payload: dict[str, Any] = self.model_dump(exclude_none=True)
+        payload.pop("metadata", None)
+        return payload
+
+
+class DataForgeOpenEnvObservation(OpenEnvObservation):
+    """OpenEnv observation model mirroring DataForge's native observation."""
+
+    visible_rows: list[dict[str, Any]] | None = None
+    detector_hints: list[str] | None = None
+    scratchpad_summary: str = ""
+    step_budget_remaining: int = 0
+    tool_usage_history: list[dict[str, Any]] = Field(default_factory=list)
+    latest_result: dict[str, Any] | None = None
+    cumulative_reward: float = 0.0
+
+
+def _to_openenv_observation(payload: dict[str, Any]) -> DataForgeOpenEnvObservation:
+    """Convert a native DataForge observation dictionary into OpenEnv shape."""
+    return DataForgeOpenEnvObservation(
+        visible_rows=payload.get("visible_rows"),
+        detector_hints=payload.get("detector_hints"),
+        scratchpad_summary=str(payload.get("scratchpad_summary", "")),
+        step_budget_remaining=int(payload.get("step_budget_remaining", 0)),
+        tool_usage_history=list(payload.get("tool_usage_history") or []),
+        latest_result=payload.get("latest_result"),
+        done=bool(payload.get("done", False)),
+        reward=payload.get("reward"),
+        cumulative_reward=float(payload.get("cumulative_reward", 0.0)),
+        metadata=dict(payload.get("metadata") or {}),
+    )
+
+
+class DataForgeOpenEnv(OpenEnvEnvironment):
+    """OpenEnv-native environment wrapper."""
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._env = DataForgeEnv()
+        self._last_observation: DataForgeOpenEnvObservation | None = None
+
+    def reset(
+        self,
+        seed: int | None = None,
+        episode_id: str | None = None,
+        **kwargs: Any,
+    ) -> DataForgeOpenEnvObservation:
+        """Reset the wrapped DataForge environment."""
+        del episode_id, kwargs
+        result = self._env.reset(seed=seed)
+        observation = _to_openenv_observation(result.observation.model_dump(mode="json"))
+        self._last_observation = observation
+        return observation
+
+    def step(
+        self,
+        action: DataForgeOpenEnvAction,
+        timeout_s: float | None = None,
+        **kwargs: Any,
+    ) -> DataForgeOpenEnvObservation:
+        """Step the wrapped DataForge environment."""
+        del timeout_s, kwargs
+        result = self._env.step(action.as_dataforge_payload())
+        observation = _to_openenv_observation(result.observation.model_dump(mode="json"))
+        self._last_observation = observation
+        return observation
+
+    def state(self) -> DataForgeOpenEnvObservation:
+        """Return the latest observation or reset lazily."""
+        if self._last_observation is None:
+            return self.reset()
+        return self._last_observation
+
+    def close(self) -> None:
+        """Close the wrapped environment."""
+        self._env.close()
+
+
+app = create_app(
+    DataForgeOpenEnv,
+    DataForgeOpenEnvAction,
+    DataForgeOpenEnvObservation,
+    env_name="dataforge-env",
+    max_concurrent_envs=64,
+)

dataforge/env/reward.py

@@ -0,0 +1,128 @@
+"""Reward engine for the DataForge RL environment.
+
+All constants and formulas are derived bit-for-bit from REWARD_DESIGN.md.
+
+Terminal score: detection_rate * 0.40 + fix_rate * 0.60 - false_positives * fp_rate
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = [
+    "DETECTION_WEIGHT",
+    "FALSE_POS_PENALTY_RATE",
+    "FIX_WEIGHT",
+    "LATE_STEP_THRESHOLD",
+    "P_FALSE_POS",
+    "P_INVALID",
+    "P_LATE_STEP",
+    "P_REINSPECT",
+    "P_WRONG_FIX",
+    "R_DIAGNOSE",
+    "R_EXPLORE",
+    "R_FIX",
+    "R_FIX_PARTIAL",
+    "R_JUSTIFY_BONUS",
+    "R_ROOT_CAUSE",
+    "R_TYPE_BONUS",
+    "SPAM_THRESHOLD",
+    "EpisodeMetrics",
+    "RewardEngine",
+]
+
+# Positive rewards
+R_DIAGNOSE: float = 0.10
+R_TYPE_BONUS: float = 0.05
+R_FIX: float = 0.15
+R_FIX_PARTIAL: float = 0.075
+R_JUSTIFY_BONUS: float = 0.05
+R_EXPLORE: float = 0.01
+R_ROOT_CAUSE: float = 0.10
+
+# Negative penalties
+P_FALSE_POS: float = -0.05
+P_WRONG_FIX: float = -0.08
+P_LATE_STEP: float = -0.02
+P_INVALID: float = -0.01
+P_REINSPECT: float = -0.01
+
+# Thresholds
+LATE_STEP_THRESHOLD: float = 0.80
+DETECTION_WEIGHT: float = 0.40
+FIX_WEIGHT: float = 0.60
+FALSE_POS_PENALTY_RATE: float = 0.05
+SPAM_THRESHOLD: float = 2.0
+
+
+@dataclass
+class EpisodeMetrics:
+    """Accumulated metrics for terminal score computation."""
+
+    found_issues: int = 0
+    total_issues: int = 0
+    fixed_issues: int = 0
+    fixable_issues: int = 0
+    false_positives: int = 0
+
+    @property
+    def total_diagnoses(self) -> int:
+        """Total diagnosis attempts (correct + incorrect)."""
+        return self.found_issues + self.false_positives
+
+
+class RewardEngine:
+    """Computes dense per-step and terminal rewards."""
+
+    def compute_terminal_score(self, metrics: EpisodeMetrics) -> float:
+        """Compute terminal score per REWARD_DESIGN.md formula."""
+        if metrics.total_issues == 0:
+            return 0.0
+        detection_rate = metrics.found_issues / metrics.total_issues
+        fix_rate = (
+            metrics.fixed_issues / metrics.fixable_issues if metrics.fixable_issues > 0 else 0.0
+        )
+        fp_rate = FALSE_POS_PENALTY_RATE
+        if (
+            metrics.total_issues > 0
+            and metrics.total_diagnoses > SPAM_THRESHOLD * metrics.total_issues
+        ):
+            fp_rate *= 2.0
+        penalty = metrics.false_positives * fp_rate
+        raw = detection_rate * DETECTION_WEIGHT + fix_rate * FIX_WEIGHT - penalty
+        return round(max(0.0, min(1.0, raw)), 4)
+
+    def compute_late_penalty(self, step: int, max_steps: int) -> float:
+        """Return P_LATE_STEP if past 80% budget, else 0.0."""
+        threshold = int(max_steps * LATE_STEP_THRESHOLD)
+        return P_LATE_STEP if step > threshold else 0.0
+
+    def compute_exploration_bonus(
+        self,
+        new_row_indices: set[int],
+        inspected_rows: set[int],
+        total_rows: int,
+        ground_truth_rows: set[int],
+        found_issue_rows: set[int],
+    ) -> float:
+        """Compute exploration bonus for newly-inspected rows."""
+        if not new_row_indices:
+            return P_REINSPECT
+        undiscovered = sum(
+            1 for r in new_row_indices if r in ground_truth_rows and r not in found_issue_rows
+        )
+        bonus = undiscovered * R_EXPLORE
+        if total_rows > 0:
+            all_inspected = inspected_rows | new_row_indices
+            coverage_ratio = len(all_inspected) / total_rows
+            bonus += len(new_row_indices) * R_EXPLORE * 0.5 * (1.0 - coverage_ratio)
+        return bonus
+
+    def diagnose_reward(self, type_match: bool) -> float:
+        """Reward for correct diagnosis."""
+        return R_DIAGNOSE + (R_TYPE_BONUS if type_match else 0.0)
+
+    def fix_reward(self, exact: bool, has_justification: bool) -> float:
+        """Reward for correct fix."""
+        reward = R_FIX if exact else R_FIX_PARTIAL
+        return reward + (R_JUSTIFY_BONUS if has_justification else 0.0)

dataforge/env/server.py

@@ -0,0 +1,176 @@
+"""FastAPI server for the DataForge RL environment.
+
+Provides OpenEnv-compatible HTTP endpoints:
+    POST /reset    — Start a new episode
+    POST /step     — Execute an action
+    GET  /state    — Return current state snapshot
+    POST /close    — No-op shutdown
+    GET  /health   — Liveness check
+    GET  /metadata — Environment metadata
+    GET  /schema   — Action/observation JSON schemas
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from threading import RLock
+from typing import Any
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import TypeAdapter
+
+from dataforge import __version__
+from dataforge.agent.tool_actions import Action
+from dataforge.env.environment import DataForgeEnv, EnvState
+from dataforge.env.observation import DataForgeObservation
+from dataforge.http.problem import problem_exception_handler
+from dataforge.observability import configure_fastapi_observability
+
+logger = logging.getLogger("dataforge.env.server")
+
+
+def _build_cors_origins() -> list[str]:
+    """Build the explicit OpenEnv CORS allowlist from the environment."""
+    raw_origins = os.environ.get("DATAFORGE_OPENENV_ORIGINS", "")
+    return [origin.strip() for origin in raw_origins.split(",") if origin.strip()]
+
+
+def _build_cors_origin_regex() -> str | None:
+    """Allow local browser development only when explicitly enabled."""
+    if os.environ.get("DATAFORGE_OPENENV_DEV") != "1":
+        return None
+    return r"^http://(?:localhost|127(?:\.\d{1,3}){3})(?::\d+)?$"
+
+
+app = FastAPI(
+    title="DataForge Environment",
+    description="OpenEnv-compatible RL environment for data-quality repair.",
+    version=__version__,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=_build_cors_origins(),
+    allow_origin_regex=_build_cors_origin_regex(),
+    allow_credentials=False,
+    allow_methods=["GET", "POST", "OPTIONS"],
+    allow_headers=["*"],
+)
+app.add_exception_handler(HTTPException, problem_exception_handler)
+configure_fastapi_observability(app, service_name="dataforge-openenv")
+
+_registry_lock = RLock()
+_default_env = DataForgeEnv()
+_sessions: dict[str, DataForgeEnv] = {}
+
+
+def _get_env(episode_id: str | None) -> DataForgeEnv:
+    """Resolve an environment by episode id, preserving legacy no-id behavior."""
+    if not episode_id:
+        return _default_env
+    with _registry_lock:
+        try:
+            return _sessions[episode_id]
+        except KeyError as exc:
+            raise HTTPException(
+                status_code=404,
+                detail={"error": "episode_not_found", "episode_id": episode_id},
+            ) from exc
+
+
+def _remember_env(env: DataForgeEnv, episode_id: str) -> None:
+    """Register a session and update the legacy default environment."""
+    global _default_env
+    with _registry_lock:
+        _sessions[episode_id] = env
+        _default_env = env
+
+
+@app.post("/reset")
+async def reset(seed: int | None = None) -> dict[str, Any]:
+    """Reset the environment for a new episode."""
+    env = DataForgeEnv()
+    result = env.reset(seed=seed)
+    episode_id = str(result.info["episode_id"])
+    _remember_env(env, episode_id)
+    return result.model_dump(mode="json")
+
+
+@app.post("/step")
+async def step(action: dict[str, Any]) -> dict[str, Any]:
+    """Execute one agent action."""
+    action_payload = dict(action)
+    raw_episode_id = action_payload.pop("episode_id", None)
+    episode_id = str(raw_episode_id) if raw_episode_id else None
+    result = _get_env(episode_id).step(action_payload)
+    return result.model_dump(mode="json")
+
+
+@app.get("/state")
+async def state(episode_id: str | None = None) -> dict[str, Any]:
+    """Return current environment state snapshot."""
+    result = _get_env(episode_id).state()
+    return result.model_dump(mode="json")
+
+
+@app.post("/close")
+async def close(request: Request, episode_id: str | None = None) -> dict[str, Any]:
+    """No-op close endpoint for OpenEnv compatibility."""
+    body_episode_id: str | None = None
+    if episode_id is None:
+        try:
+            payload = await request.json()
+        except Exception:
+            payload = None
+        if isinstance(payload, dict) and payload.get("episode_id"):
+            body_episode_id = str(payload["episode_id"])
+
+    target_episode_id = episode_id or body_episode_id
+    env = _get_env(target_episode_id)
+    env.close()
+    if target_episode_id:
+        with _registry_lock:
+            _sessions.pop(target_episode_id, None)
+    return {"status": "closed", "episode_id": target_episode_id}
+
+
+@app.get("/health")
+async def health() -> dict[str, Any]:
+    """Liveness check."""
+    return {"status": "healthy", "environment": "dataforge-env"}
+
+
+@app.get("/metadata")
+async def metadata() -> dict[str, Any]:
+    """Environment metadata for OpenEnv discovery."""
+    return {
+        "name": "dataforge-env",
+        "version": __version__,
+        "description": (
+            "DataForge RL Environment — agents learn to detect, diagnose, "
+            "and repair data-quality issues in tabular datasets."
+        ),
+        "action_types": [
+            "INSPECT_ROWS",
+            "SQL_QUERY",
+            "STAT_TEST",
+            "PATTERN_MATCH",
+            "HYPOTHESIS",
+            "ROOT_CAUSE",
+            "DIAGNOSE",
+            "FIX",
+        ],
+    }
+
+
+@app.get("/schema")
+async def schema() -> dict[str, Any]:
+    """Return JSON schemas for action and observation models."""
+    action_adapter: TypeAdapter[Action] = TypeAdapter(Action)
+    return {
+        "action": action_adapter.json_schema(),
+        "observation": DataForgeObservation.model_json_schema(),
+        "state": EnvState.model_json_schema(),
+    }

dataforge/evaluation_contract.py

@@ -0,0 +1,76 @@
+"""Public evaluation evidence models for DataForge repair releases."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+InferabilityLabel = Literal[
+    "deterministic_normalization",
+    "context_derivable",
+    "external_reference_required",
+    "not_inferable_from_prompt",
+]
+PROMOTION_SLICE: InferabilityLabel = "deterministic_normalization"
+ABSTENTION_SLICES = frozenset({"external_reference_required", "not_inferable_from_prompt"})
+AUXILIARY_SLICES = frozenset({"context_derivable"})
+PromotionStatus = Literal[
+    "diagnostic_only",
+    "diagnostic_promoted",
+    "quality_improved_verified",
+    "public_quality_milestone",
+    "rejected",
+]
+
+
+class EvaluationTaskV2(BaseModel):
+    """One auditable, source-stable model grading task.
+
+    Ground truth is retained for local grading but excluded from normal JSON
+    serialization so prompts and public reports cannot accidentally leak labels.
+    """
+
+    schema_version: Literal["evaluation_task_v2"] = "evaluation_task_v2"
+    task_id: str = Field(min_length=1)
+    prompt_hash: str = Field(min_length=64, max_length=64)
+    dataset_sha: str = Field(min_length=1)
+    split_id: str = Field(min_length=1)
+    inferability: InferabilityLabel
+    prompt: dict[str, Any]
+    allowed_columns: list[str] = Field(min_length=1)
+    valid_rows: list[int] = Field(min_length=1)
+    provenance: dict[str, Any]
+    hidden_ground_truth: list[dict[str, Any]] = Field(default_factory=list, exclude=True)
+
+    model_config = {"frozen": True}
+
+
+class ReleaseEvidenceV2(BaseModel):
+    """Serializable release-gate evidence for model and benchmark promotion."""
+
+    schema_version: Literal["release_evidence_v2"] = "release_evidence_v2"
+    model_repo: str = Field(min_length=1)
+    model_sha: str = Field(min_length=1)
+    dataset_repo: str = Field(min_length=1)
+    dataset_sha: str = Field(min_length=1)
+    strict_macro_f1: float = Field(ge=0.0, le=1.0)
+    canonicalized_macro_f1: float = Field(ge=0.0, le=1.0)
+    parse_success_rate: float = Field(ge=0.0, le=1.0)
+    schema_case_error_count: int = Field(ge=0)
+    promotion_slice: InferabilityLabel = PROMOTION_SLICE
+    slice_scores: dict[InferabilityLabel, dict[str, float | int]] = Field(default_factory=dict)
+    inferability_slice_scores: dict[InferabilityLabel, float] = Field(default_factory=dict)
+    package_versions: dict[str, str] = Field(default_factory=dict)
+    promotion_status: PromotionStatus
+    gate_failures: list[str] = Field(default_factory=list)
+
+    model_config = {"frozen": True}
+
+
+def prompt_sha256(prompt: dict[str, Any]) -> str:
+    """Hash a prompt payload with stable JSON serialization."""
+    encoded = json.dumps(prompt, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()

dataforge/fixtures/hospital_10rows.csv

@@ -0,0 +1,11 @@
+provider_number,hospital_name,city,state,zip_code,phone_number,rating,mortality_rate,readmission_rate,er_wait_time
+PRV001,General Hospital,Springfield,IL,62701,2175550101,4.2,0.023,0.145,28
+PRV002,St. Mary Medical Center,Chicago,IL,60601,3125550202,3.8,0.031,0.162,35
+PRV001,Springfield Medical,Springfield,IL,62701,2175550303,4.5,0.019,0.138,22
+PRV003,Mercy Hospital,Peoria,IL,61602,3095550404,3.5,0.028,0.158,31
+PRV004,Northwestern Memorial,Chicago,IL,60611,not available,4.1,0.025,0.149,26
+PRV005,Rush University MC,Chicago,IL,60612,3125550606,45.0,0.022,0.141,29
+PRV006,Advocate Christ,Oak Lawn,IL,60453,7085550707,3.9,0.027,0.155,33
+PRV007,Loyola University MC,Maywood,IL,60153,7085550808,4.3,0.020,0.142,25
+PRV008,Presence St. Joseph,Joliet,IL,60435,8155550909,4.0,0.026,0.151,30
+PRV009,Edward Hospital,Naperville,IL,60540,6305551010,3.7,0.029,0.160,34

dataforge/fixtures/hospital_schema.yaml

@@ -0,0 +1,17 @@
+# Hospital dataset schema for DataForge profile command.
+
+columns:
+  provider_number: str
+  hospital_name: str
+  city: str
+  state: str
+  zip_code: str
+  phone_number: str
+  rating: float
+  mortality_rate: float
+  readmission_rate: float
+  er_wait_time: int
+
+functional_dependencies:
+  - determinant: [provider_number]
+    dependent: hospital_name

dataforge/http/__init__.py

@@ -0,0 +1 @@
+"""HTTP helpers shared by DataForge backend surfaces."""