dataforge-07 0.1.0__py3-none-any.whl

dataforge/__init__.py

@@ -0,0 +1,204 @@
+"""DataForge public package.
+
+The root package is the stable facade for integration surfaces. Symbols are
+resolved lazily so importing :mod:`dataforge` does not eagerly import pandas,
+FastAPI-facing helpers, or the SMT stack.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from dataforge.cli.common import load_schema, schema_from_mapping
+    from dataforge.detectors import Issue, Schema, Severity, run_all_detectors
+    from dataforge.engine.repair import (
+        CandidateFix,
+        CandidateRepair,
+        ProofObligation,
+        RepairFailure,
+        RepairPipelineRequest,
+        RepairPipelineResult,
+        RepairReceipt,
+        RootCause,
+        VerifiedFix,
+        run_repair_pipeline,
+    )
+    from dataforge.integrations.dbt import schema_from_dbt_artifacts, schema_from_dbt_manifest
+    from dataforge.repair_contract import CONTRACT_VERSION
+    from dataforge.repairers import ProposedFix
+    from dataforge.safety import SafetyContext, SafetyFilter, SafetyResult, SafetyVerdict
+    from dataforge.schema_inference import (
+        ConstraintCandidate,
+        ConstraintReviewArtifact,
+        ConstraintReviewError,
+        ReviewedConstraintCandidate,
+        SchemaInferenceResult,
+        build_constraint_review_artifact,
+        dump_constraint_review_artifact,
+        infer_schema,
+        load_constraint_review_artifact,
+    )
+    from dataforge.stores import (
+        DuckDBStore,
+        PatchPlan,
+        TableStoreError,
+        TableStoreRepairResult,
+        is_table_store_uri,
+        run_table_store_repair,
+        store_from_uri,
+    )
+    from dataforge.table import read_csv
+    from dataforge.transactions.log import (
+        TransactionAuditReport,
+        TransactionAuditVerdict,
+        TransactionLogError,
+        verify_transaction_log,
+    )
+    from dataforge.transactions.revert import TransactionRevertError, revert_transaction
+    from dataforge.transactions.txn import CellFix, RepairTransaction
+    from dataforge.verifier import (
+        ConstraintIR,
+        SMTVerifier,
+        VerificationResult,
+        VerificationVerdict,
+        constraint_ir_from_schema,
+    )
+
+__all__ = [
+    "CONTRACT_VERSION",
+    "CandidateFix",
+    "CandidateRepair",
+    "CellFix",
+    "ConstraintCandidate",
+    "ConstraintReviewArtifact",
+    "ConstraintReviewError",
+    "ConstraintIR",
+    "DuckDBStore",
+    "Issue",
+    "PatchPlan",
+    "ProposedFix",
+    "ProofObligation",
+    "RepairFailure",
+    "RepairPipelineRequest",
+    "RepairPipelineResult",
+    "RepairReceipt",
+    "RepairTransaction",
+    "RootCause",
+    "ReviewedConstraintCandidate",
+    "SMTVerifier",
+    "SafetyContext",
+    "SafetyFilter",
+    "SafetyResult",
+    "SafetyVerdict",
+    "Schema",
+    "SchemaInferenceResult",
+    "Severity",
+    "TransactionAuditReport",
+    "TransactionAuditVerdict",
+    "TransactionLogError",
+    "TransactionRevertError",
+    "TableStoreError",
+    "TableStoreRepairResult",
+    "VerificationResult",
+    "VerificationVerdict",
+    "VerifiedFix",
+    "__version__",
+    "load_schema",
+    "build_constraint_review_artifact",
+    "constraint_ir_from_schema",
+    "dump_constraint_review_artifact",
+    "load_constraint_review_artifact",
+    "read_csv",
+    "revert_transaction",
+    "run_all_detectors",
+    "run_repair_pipeline",
+    "schema_from_mapping",
+    "schema_from_dbt_artifacts",
+    "schema_from_dbt_manifest",
+    "infer_schema",
+    "is_table_store_uri",
+    "run_table_store_repair",
+    "store_from_uri",
+    "verify_transaction_log",
+]
+
+__version__ = "0.1.0"
+
+_PUBLIC_EXPORTS: dict[str, tuple[str, str]] = {
+    "CONTRACT_VERSION": ("dataforge.repair_contract", "CONTRACT_VERSION"),
+    "CandidateFix": ("dataforge.engine.repair", "CandidateFix"),
+    "CandidateRepair": ("dataforge.engine.repair", "CandidateRepair"),
+    "CellFix": ("dataforge.transactions.txn", "CellFix"),
+    "ConstraintCandidate": ("dataforge.schema_inference", "ConstraintCandidate"),
+    "ConstraintReviewArtifact": ("dataforge.schema_inference", "ConstraintReviewArtifact"),
+    "ConstraintReviewError": ("dataforge.schema_inference", "ConstraintReviewError"),
+    "ConstraintIR": ("dataforge.verifier", "ConstraintIR"),
+    "DuckDBStore": ("dataforge.stores", "DuckDBStore"),
+    "Issue": ("dataforge.detectors", "Issue"),
+    "ProposedFix": ("dataforge.repairers", "ProposedFix"),
+    "ProofObligation": ("dataforge.engine.repair", "ProofObligation"),
+    "PatchPlan": ("dataforge.stores", "PatchPlan"),
+    "RepairFailure": ("dataforge.engine.repair", "RepairFailure"),
+    "RepairPipelineRequest": ("dataforge.engine.repair", "RepairPipelineRequest"),
+    "RepairPipelineResult": ("dataforge.engine.repair", "RepairPipelineResult"),
+    "RepairReceipt": ("dataforge.engine.repair", "RepairReceipt"),
+    "RepairTransaction": ("dataforge.transactions.txn", "RepairTransaction"),
+    "RootCause": ("dataforge.engine.repair", "RootCause"),
+    "ReviewedConstraintCandidate": ("dataforge.schema_inference", "ReviewedConstraintCandidate"),
+    "SMTVerifier": ("dataforge.verifier", "SMTVerifier"),
+    "SafetyContext": ("dataforge.safety", "SafetyContext"),
+    "SafetyFilter": ("dataforge.safety", "SafetyFilter"),
+    "SafetyResult": ("dataforge.safety", "SafetyResult"),
+    "SafetyVerdict": ("dataforge.safety", "SafetyVerdict"),
+    "Schema": ("dataforge.detectors", "Schema"),
+    "SchemaInferenceResult": ("dataforge.schema_inference", "SchemaInferenceResult"),
+    "Severity": ("dataforge.detectors", "Severity"),
+    "TransactionAuditReport": ("dataforge.transactions.log", "TransactionAuditReport"),
+    "TransactionAuditVerdict": ("dataforge.transactions.log", "TransactionAuditVerdict"),
+    "TransactionLogError": ("dataforge.transactions.log", "TransactionLogError"),
+    "TransactionRevertError": ("dataforge.transactions.revert", "TransactionRevertError"),
+    "TableStoreError": ("dataforge.stores", "TableStoreError"),
+    "TableStoreRepairResult": ("dataforge.stores", "TableStoreRepairResult"),
+    "VerificationResult": ("dataforge.verifier", "VerificationResult"),
+    "VerificationVerdict": ("dataforge.verifier", "VerificationVerdict"),
+    "VerifiedFix": ("dataforge.engine.repair", "VerifiedFix"),
+    "load_schema": ("dataforge.cli.common", "load_schema"),
+    "build_constraint_review_artifact": (
+        "dataforge.schema_inference",
+        "build_constraint_review_artifact",
+    ),
+    "constraint_ir_from_schema": ("dataforge.verifier", "constraint_ir_from_schema"),
+    "dump_constraint_review_artifact": (
+        "dataforge.schema_inference",
+        "dump_constraint_review_artifact",
+    ),
+    "load_constraint_review_artifact": (
+        "dataforge.schema_inference",
+        "load_constraint_review_artifact",
+    ),
+    "read_csv": ("dataforge.table", "read_csv"),
+    "revert_transaction": ("dataforge.transactions.revert", "revert_transaction"),
+    "run_all_detectors": ("dataforge.detectors", "run_all_detectors"),
+    "run_repair_pipeline": ("dataforge.engine.repair", "run_repair_pipeline"),
+    "schema_from_mapping": ("dataforge.cli.common", "schema_from_mapping"),
+    "schema_from_dbt_artifacts": ("dataforge.integrations.dbt", "schema_from_dbt_artifacts"),
+    "schema_from_dbt_manifest": ("dataforge.integrations.dbt", "schema_from_dbt_manifest"),
+    "infer_schema": ("dataforge.schema_inference", "infer_schema"),
+    "is_table_store_uri": ("dataforge.stores", "is_table_store_uri"),
+    "run_table_store_repair": ("dataforge.stores", "run_table_store_repair"),
+    "store_from_uri": ("dataforge.stores", "store_from_uri"),
+    "verify_transaction_log": ("dataforge.transactions.log", "verify_transaction_log"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    """Resolve public facade exports on first use."""
+    try:
+        module_name, attribute_name = _PUBLIC_EXPORTS[name]
+    except KeyError as exc:
+        raise AttributeError(name) from exc
+    value = getattr(import_module(module_name), attribute_name)
+    globals()[name] = value
+    return value

dataforge/__main__.py

@@ -0,0 +1,5 @@
+"""Enable ``python -m dataforge`` invocation."""
+
+from dataforge.cli import app
+
+app()

dataforge/agent/__init__.py

@@ -0,0 +1,16 @@
+"""DataForge agent package — typed tool-use actions and scratchpad.
+
+Public API:
+    parse_action — Parse raw dict into typed Action model.
+    Action       — Discriminated union of all action types.
+    Scratchpad   — In-episode hypothesis tracker.
+"""
+
+from dataforge.agent.scratchpad import Scratchpad
+from dataforge.agent.tool_actions import Action, parse_action
+
+__all__ = [
+    "Action",
+    "Scratchpad",
+    "parse_action",
+]

dataforge/agent/providers.py

@@ -0,0 +1,259 @@
+"""Multi-provider LLM client for DataForge.
+
+Reads ``DATAFORGE_LLM_PROVIDER`` from the environment and dispatches to the
+matching provider.  Week 1 implements **groq** and **gemini** only; other
+providers raise ``NotImplementedError``.
+
+No LLM calls are made by detectors — this module is for the agent loop
+(Week 2+) and is stubbed here to establish the interface.
+
+The interface is:
+    ``async def complete(messages, model, temperature) -> str``
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Literal, TypedDict
+
+import httpx
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+
+# ── Message type ──────────────────────────────────────────────────────────
+
+
+class Message(TypedDict):
+    """A single chat message.
+
+    Args:
+        role: The speaker role — ``"system"``, ``"user"``, or ``"assistant"``.
+        content: The text content of the message.
+    """
+
+    role: Literal["system", "user", "assistant"]
+    content: str
+
+
+# ── Exceptions ────────────────────────────────────────────────────────────
+
+
+class ProviderError(Exception):
+    """Raised when an LLM provider call fails after retries.
+
+    Args:
+        provider: The provider name that failed.
+        message: Description of the failure.
+    """
+
+    def __init__(self, provider: str, message: str) -> None:
+        self.provider = provider
+        super().__init__(f"[{provider}] {message}")
+
+
+# ── Provider dispatch ─────────────────────────────────────────────────────
+
+_SUPPORTED_PROVIDERS = frozenset({"groq", "gemini", "cerebras", "openrouter", "hf", "cloudflare"})
+
+
+def get_provider_name() -> str:
+    """Read the active provider from the environment.
+
+    Returns:
+        The lowercased provider name from ``DATAFORGE_LLM_PROVIDER``.
+        When no explicit provider is configured, prefer a provider whose
+        credential is present in the environment.
+
+    Example:
+        >>> import os
+        >>> os.environ["DATAFORGE_LLM_PROVIDER"] = "gemini"
+        >>> get_provider_name()
+        'gemini'
+    """
+    configured = os.environ.get("DATAFORGE_LLM_PROVIDER")
+    if configured:
+        return configured.lower()
+    if os.environ.get("GROQ_API_KEY"):
+        return "groq"
+    if os.environ.get("GEMINI_API_KEY"):
+        return "gemini"
+    return "groq"
+
+
+async def complete(
+    messages: list[Message],
+    *,
+    model: str | None = None,
+    temperature: float = 0.0,
+) -> str:
+    """Send a chat completion request to the active LLM provider.
+
+    Args:
+        messages: List of chat messages forming the conversation.
+        model: Optional model override. If None, uses the provider default.
+        temperature: Sampling temperature (0.0 = deterministic).
+
+    Returns:
+        The assistant's response text.
+
+    Raises:
+        NotImplementedError: If the provider is not yet implemented.
+        ProviderError: If the API call fails after retries.
+
+    Example:
+        >>> import asyncio
+        >>> msgs = [{"role": "user", "content": "What is 2+2?"}]
+        >>> # result = asyncio.run(complete(msgs))  # requires API key
+    """
+    provider = get_provider_name()
+
+    if provider == "groq":
+        return await _complete_groq(messages, model=model, temperature=temperature)
+    if provider == "gemini":
+        return await _complete_gemini(messages, model=model, temperature=temperature)
+
+    if provider in _SUPPORTED_PROVIDERS:
+        raise NotImplementedError(
+            f"Provider '{provider}' is planned but not yet implemented. "
+            f"Use 'groq' or 'gemini' for Week 1."
+        )
+
+    raise NotImplementedError(
+        f"Unknown provider '{provider}'. Supported: {sorted(_SUPPORTED_PROVIDERS)}"
+    )
+
+
+# ── Groq provider ────────────────────────────────────────────────────────
+
+_GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
+_GROQ_DEFAULT_MODEL = "llama-3.1-70b-versatile"
+
+
+@retry(
+    retry=retry_if_exception_type(httpx.HTTPStatusError),
+    wait=wait_exponential(multiplier=1, min=1, max=30),
+    stop=stop_after_attempt(3),
+    reraise=True,
+)
+async def _complete_groq(
+    messages: list[Message],
+    *,
+    model: str | None = None,
+    temperature: float = 0.0,
+) -> str:
+    """Call Groq's OpenAI-compatible chat completions API.
+
+    Args:
+        messages: Chat messages.
+        model: Model name (defaults to llama-3.1-70b-versatile).
+        temperature: Sampling temperature.
+
+    Returns:
+        The assistant's response text.
+
+    Raises:
+        ProviderError: If the response is malformed.
+    """
+    api_key = os.environ.get("GROQ_API_KEY", "")
+    if not api_key:
+        raise ProviderError("groq", "GROQ_API_KEY environment variable not set")
+
+    payload = {
+        "model": model or _GROQ_DEFAULT_MODEL,
+        "messages": [dict(m) for m in messages],
+        "temperature": temperature,
+    }
+
+    async with httpx.AsyncClient(timeout=60.0) as client:
+        response = await client.post(
+            _GROQ_URL,
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            },
+        )
+        response.raise_for_status()
+
+    data = response.json()
+    try:
+        return str(data["choices"][0]["message"]["content"])
+    except (KeyError, IndexError) as exc:
+        raise ProviderError("groq", f"Unexpected response format: {data}") from exc
+
+
+# ── Gemini provider ──────────────────────────────────────────────────────
+
+_GEMINI_URL_TEMPLATE = (
+    "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
+)
+_GEMINI_DEFAULT_MODEL = "gemini-2.0-flash"
+
+
+@retry(
+    retry=retry_if_exception_type(httpx.HTTPStatusError),
+    wait=wait_exponential(multiplier=1, min=1, max=30),
+    stop=stop_after_attempt(3),
+    reraise=True,
+)
+async def _complete_gemini(
+    messages: list[Message],
+    *,
+    model: str | None = None,
+    temperature: float = 0.0,
+) -> str:
+    """Call Google's Gemini generativeLanguage API.
+
+    Args:
+        messages: Chat messages (converted to Gemini's content format).
+        model: Model name (defaults to gemini-2.0-flash).
+        temperature: Sampling temperature.
+
+    Returns:
+        The assistant's response text.
+
+    Raises:
+        ProviderError: If the response is malformed.
+    """
+    api_key = os.environ.get("GEMINI_API_KEY", "")
+    if not api_key:
+        raise ProviderError("gemini", "GEMINI_API_KEY environment variable not set")
+
+    model_name = model or _GEMINI_DEFAULT_MODEL
+    url = _GEMINI_URL_TEMPLATE.format(model=model_name)
+
+    # Convert OpenAI-style messages to Gemini format.
+    contents: list[dict[str, object]] = []
+    system_instruction: str | None = None
+    for msg in messages:
+        if msg["role"] == "system":
+            system_instruction = msg["content"]
+        else:
+            role = "user" if msg["role"] == "user" else "model"
+            contents.append(
+                {
+                    "role": role,
+                    "parts": [{"text": msg["content"]}],
+                }
+            )
+
+    payload: dict[str, object] = {
+        "contents": contents,
+        "generationConfig": {"temperature": temperature},
+    }
+    if system_instruction:
+        payload["systemInstruction"] = {"parts": [{"text": system_instruction}]}
+
+    async with httpx.AsyncClient(timeout=60.0) as client:
+        response = await client.post(
+            url,
+            json=payload,
+            params={"key": api_key},
+            headers={"Content-Type": "application/json"},
+        )
+        response.raise_for_status()
+
+    data = response.json()
+    try:
+        return str(data["candidates"][0]["content"]["parts"][0]["text"])
+    except (KeyError, IndexError) as exc:
+        raise ProviderError("gemini", f"Unexpected response format: {data}") from exc

dataforge/agent/scratchpad.py

@@ -0,0 +1,183 @@
+"""In-episode hypothesis and issue tracker for the DataForge RL agent.
+
+The scratchpad is a mutable, episode-scoped data structure that the agent
+uses to record hypotheses, confirmed issues, and dead ends. The environment
+exposes a compact summary of the scratchpad in each observation, enabling
+the agent to reason about its investigation history without direct access
+to the underlying data structure.
+
+Example::
+
+    >>> from dataforge.agent.scratchpad import Scratchpad
+    >>> pad = Scratchpad()
+    >>> pad.add_hypothesis("Rating column has decimal shift", [5], ["rating"], "decimal_shift")
+    >>> pad.confirm_issue(5, "rating", "decimal_shift")
+    >>> pad.summary()
+    'Hypotheses: 1 (0 pending). Confirmed: 1. Dead ends: 0.'
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+__all__ = [
+    "ConfirmedIssue",
+    "DeadEnd",
+    "HypothesisRecord",
+    "Scratchpad",
+]
+
+
+@dataclass(frozen=True)
+class HypothesisRecord:
+    """A recorded hypothesis about a data-quality root cause.
+
+    Args:
+        claim: Textual description of the hypothesis.
+        affected_rows: Row indices the hypothesis covers.
+        affected_columns: Column names the hypothesis covers.
+        root_cause_type: Detector-vocabulary root cause type.
+        confirmed: Whether the hypothesis was confirmed by ground truth.
+    """
+
+    claim: str
+    affected_rows: tuple[int, ...]
+    affected_columns: tuple[str, ...]
+    root_cause_type: str
+    confirmed: bool = False
+
+
+@dataclass(frozen=True)
+class ConfirmedIssue:
+    """A confirmed data-quality issue at a specific location.
+
+    Args:
+        row: Zero-indexed row number.
+        column: Column name.
+        issue_type: Issue type classification.
+    """
+
+    row: int
+    column: str
+    issue_type: str
+
+
+@dataclass(frozen=True)
+class DeadEnd:
+    """A recorded dead end — an investigation path that yielded nothing.
+
+    Args:
+        description: What was tried and why it failed.
+        step_number: Step at which the dead end was recorded.
+    """
+
+    description: str
+    step_number: int
+
+
+@dataclass
+class Scratchpad:
+    """Mutable in-episode tracker for hypotheses, confirmed issues, and dead ends.
+
+    Reset at the start of each episode. The ``summary()`` method produces a
+    compact string for inclusion in agent observations.
+
+    Example::
+
+        >>> pad = Scratchpad()
+        >>> pad.add_hypothesis("Decimal shift in rating", [5], ["rating"], "decimal_shift")
+        >>> len(pad.hypotheses)
+        1
+    """
+
+    hypotheses: list[HypothesisRecord] = field(default_factory=list)
+    confirmed_issues: list[ConfirmedIssue] = field(default_factory=list)
+    dead_ends: list[DeadEnd] = field(default_factory=list)
+
+    def add_hypothesis(
+        self,
+        claim: str,
+        affected_rows: list[int],
+        affected_columns: list[str],
+        root_cause_type: str,
+    ) -> HypothesisRecord:
+        """Record a new hypothesis.
+
+        Args:
+            claim: Textual description of the hypothesis.
+            affected_rows: Row indices the hypothesis covers.
+            affected_columns: Column names the hypothesis covers.
+            root_cause_type: Detector-vocabulary root cause type.
+
+        Returns:
+            The recorded hypothesis.
+        """
+        record = HypothesisRecord(
+            claim=claim,
+            affected_rows=tuple(affected_rows),
+            affected_columns=tuple(affected_columns),
+            root_cause_type=root_cause_type,
+        )
+        self.hypotheses.append(record)
+        return record
+
+    def confirm_hypothesis(self, index: int) -> None:
+        """Mark a hypothesis as confirmed.
+
+        Args:
+            index: Index into the ``hypotheses`` list.
+
+        Raises:
+            IndexError: If the index is out of range.
+        """
+        old = self.hypotheses[index]
+        self.hypotheses[index] = HypothesisRecord(
+            claim=old.claim,
+            affected_rows=old.affected_rows,
+            affected_columns=old.affected_columns,
+            root_cause_type=old.root_cause_type,
+            confirmed=True,
+        )
+
+    def confirm_issue(self, row: int, column: str, issue_type: str) -> None:
+        """Record a confirmed issue.
+
+        Args:
+            row: Zero-indexed row number.
+            column: Column name.
+            issue_type: Issue type classification.
+        """
+        self.confirmed_issues.append(ConfirmedIssue(row=row, column=column, issue_type=issue_type))
+
+    def add_dead_end(self, description: str, step_number: int) -> None:
+        """Record a dead end.
+
+        Args:
+            description: What was tried and why it failed.
+            step_number: Step at which the dead end was recorded.
+        """
+        self.dead_ends.append(DeadEnd(description=description, step_number=step_number))
+
+    def reset(self) -> None:
+        """Clear all tracked state for a new episode."""
+        self.hypotheses.clear()
+        self.confirmed_issues.clear()
+        self.dead_ends.clear()
+
+    def summary(self) -> str:
+        """Produce a compact summary string for observation embedding.
+
+        Returns:
+            A one-line summary of scratchpad state.
+
+        Example::
+
+            >>> Scratchpad().summary()
+            'Hypotheses: 0 (0 pending). Confirmed: 0. Dead ends: 0.'
+        """
+        pending = sum(1 for h in self.hypotheses if not h.confirmed)
+        return (
+            f"Hypotheses: {len(self.hypotheses)} ({pending} pending). "
+            f"Confirmed: {len(self.confirmed_issues)}. "
+            f"Dead ends: {len(self.dead_ends)}."
+        )