"""Adaptive, auditable BuffData optimization pipeline."""

from __future__ import annotations

import hashlib

import json

import orjson

import os

import asyncio

from pathlib import Path

from typing import Any, List, Optional

from buffdata.engine.client import LLMClient, create_llm_client

from buffdata.engine.limiter import AsyncRateLimiter

from buffdata.engine.profiler import DatasetProfiler, representative_sample

from buffdata.engine.validator import DatasetValidator

from buffdata.models.formats import read_dataset, write_dataset_atomic

from buffdata.observability import ObservabilityRegistry, disabled_registry

from buffdata.security.permissions import restrict_to_owner

from buffdata.models.schemas import (

ClassificationMode,

DatasetItem,

OptimizationRunResult,

PipelineConfig,

)

from buffdata.optimizers.classifier import ClassificationSchema, DatasetClassifier

from buffdata.optimizers.dedup import Deduplicator

from buffdata.optimizers.refiner import DataRefiner

from buffdata.optimizers.scorer import FastRuleFilter, QualityScorer

from buffdata.optimizers.scrubber import PIIScrubber

from buffdata.report.generator import ReportGenerator

STAGES = ("validate", "pii", "profile", "dedup", "score_refine", "filter", "classify")

class OptimizationPipeline:

def __init__(self, config: PipelineConfig, client: Optional[LLMClient] = None):

self.config = config

self.client = client or create_llm_client(

provider=config.provider,

model=config.model,

allow_mock=False,

base_url=config.base_url,

network_policy=config.network_policy,

)

self.model = config.model or self.client.default_model

self.limiter = AsyncRateLimiter(

max_rpm=config.max_rpm,

concurrency=config.concurrency,

)

self._observability = (

ObservabilityRegistry(enabled=True) if config.observability else disabled_registry()

)

self._strict_labeled: Optional[bool] = None

@staticmethod

def artifact_paths(output_path: Path) -> tuple[Path, Path, Path]:

rejected = output_path.with_name(f"{output_path.stem}.rejected.jsonl")

report = output_path.with_name(f"{output_path.stem}.report.json")

checkpoint = output_path.with_name(f".{output_path.stem}.checkpoint.json")

return rejected, report, checkpoint

def _fingerprint(self, input_path: Path) -> str:

digest = hashlib.sha256()

if input_path.is_dir():

for child in sorted(path for path in input_path.rglob("*") if path.is_file()):

digest.update(child.relative_to(input_path).as_posix().encode("utf-8"))

with child.open("rb") as handle:

for block in iter(lambda: handle.read(1024 * 1024), b""):

digest.update(block)

else:

with input_path.open("rb") as handle:

for block in iter(lambda: handle.read(1024 * 1024), b""):

digest.update(block)

digest.update(self.config.model_dump_json(exclude_none=False).encode("utf-8"))

return digest.hexdigest()

@staticmethod

def _reject(item: DatasetItem, stage: str, reason: str) -> None:

item.metadata["rejection"] = {"stage": stage, "reason": reason}

@staticmethod

def _atomic_json(path: Path, data: dict[str, Any]) -> None:

path.parent.mkdir(parents=True, exist_ok=True)

temporary = path.with_name(f".{path.name}.tmp")

temporary.write_bytes(orjson.dumps(data, option=orjson.OPT_INDENT_2))

restrict_to_owner(temporary)

os.replace(temporary, path)

def _is_strict_labeled(self, items: List[DatasetItem]) -> bool:

"""Cached check for strict labeled mode - avoids repeated all() calls."""

if self._strict_labeled is None:

self._strict_labeled = (

self.config.accuracy_contract == "strict"

and bool(items)

and all(item.labels is not None for item in items)

)

return self._strict_labeled

async def _save_checkpoint(

self,

path: Optional[Path],

fingerprint: Optional[str],

completed: List[str],

accepted: List[DatasetItem],

rejected: List[DatasetItem],

profile: Any,

metrics: dict[str, Any],

) -> None:

if path is None or fingerprint is None:

return

self._atomic_json(

path,

{

"fingerprint": fingerprint,

"completed_stages": completed,

"accepted": [item.to_dict() for item in accepted],

"rejected": [item.to_dict() for item in rejected],

"profile": profile.model_dump(mode="json") if profile else None,

"metrics": metrics,

},

)

@staticmethod

def _load_checkpoint(path: Path, fingerprint: str):

if not path.exists():

return None

data = json.loads(path.read_text(encoding="utf-8"))

if data.get("fingerprint") != fingerprint:

raise ValueError(

f"Checkpoint {path} belongs to a different input or configuration. "

"Remove it or restore the matching run."

)

return data

async def run_file(self, input_path: Path | str, output_path: Path | str) -> OptimizationRunResult:

input_file = Path(input_path)

output_file = Path(output_path)

rejected_path, report_path, checkpoint_path = self.artifact_paths(output_file)

fingerprint = self._fingerprint(input_file)

checkpoint = self._load_checkpoint(checkpoint_path, fingerprint)

items = read_dataset(input_file)

return await self.run(

items,

output_path=output_file,

rejected_path=rejected_path,

report_path=report_path,

checkpoint_path=checkpoint_path,

fingerprint=fingerprint,

resume_state=checkpoint,

)

async def run(

self,

items: List[DatasetItem],

*,

output_path: Optional[Path] = None,

rejected_path: Optional[Path] = None,

report_path: Optional[Path] = None,

checkpoint_path: Optional[Path] = None,

fingerprint: Optional[str] = None,

resume_state: Optional[dict[str, Any]] = None,

) -> OptimizationRunResult:

accepted = list(items)

rejected: List[DatasetItem] = []

completed: List[str] = []

profile = None

metrics: dict[str, Any] = {"input_records": len(items), "stages": {}}

if resume_state:

accepted = [DatasetItem.from_dict(row) for row in resume_state.get("accepted", [])]

rejected = [DatasetItem.from_dict(row) for row in resume_state.get("rejected", [])]

completed = list(resume_state.get("completed_stages", []))

metrics = resume_state.get("metrics", metrics)

profile_data = resume_state.get("profile")

if profile_data:

from buffdata.models.schemas import DatasetProfile

profile = DatasetProfile.model_validate(profile_data)

if "validate" not in completed:

rejected_before = len(rejected)

with self._observability.stage_span("validate"):

# Structural validation only - FastRuleFilter is applied later in score_refine stage

valid = []

invalid = []

for item in accepted:

errors = DatasetValidator.validate_item(item)

if errors:

self._reject(item, "validate", "; ".join(errors))

invalid.append(item)

else:

valid.append(item)

accepted = valid

rejected.extend(invalid)

metrics["stages"]["validate"] = {"accepted": len(accepted), "rejected": len(invalid)}

self._observability.record_stage_outcome(

"validate", accepted=len(accepted), rejected=len(invalid)

)

completed.append("validate")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "pii" not in completed:

with self._observability.stage_span("pii"):

strict_labeled = self._is_strict_labeled(accepted)

labeled_classification = bool(accepted) and all(item.labels is not None for item in accepted)

classification_mode = (

"off"

if strict_labeled

else self.config.classification_pii_mode if labeled_classification else "all"

)

pii_enabled = self.config.scrub_pii and classification_mode != "off"

if pii_enabled:

entities = (

list(PIIScrubber.IDENTIFIER_ENTITIES)

if labeled_classification and classification_mode == "identifiers"

else None

)

scrubber = PIIScrubber(entities=entities)

accepted = await asyncio.to_thread(scrubber.scrub_batch, accepted)

metrics["stages"]["pii"] = {

"enabled": pii_enabled,

"policy": classification_mode,

"redacted_records": sum(

bool(item.metadata.get("pii", {}).get("entity_counts")) for item in accepted

),

}

self._observability.record_stage_outcome("pii", accepted=len(accepted))

completed.append("pii")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "profile" not in completed:

with self._observability.stage_span("profile"):

profile = await DatasetProfiler(

client=self.client,

model=self.model,

sample_size=self.config.classification_sample_size,

).profile(accepted)

metrics["stages"]["profile"] = {

"classification_applicable": profile.classification_applicable,

"confidence": profile.confidence,

"task_type": profile.task_type.value if profile.task_type else None,

}

self._observability.record_stage_outcome("profile", accepted=len(accepted))

completed.append("profile")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "dedup" not in completed:

with self._observability.stage_span("dedup"):

deduper = Deduplicator(client=self.client)

method = "off" if self._is_strict_labeled(accepted) else self.config.dedup_method

if method == "off":

duplicates = []

else:

if method == "auto":

method = (

"exact"

if accepted and all(item.labels is not None for item in accepted)

else "minhash"

)

if method == "exact":

accepted, duplicates = await asyncio.to_thread(deduper.deduplicate_exact, accepted)

elif method == "semantic-local":

accepted, duplicates = await asyncio.to_thread(

deduper.deduplicate_semantic_local,

accepted,

threshold=self.config.dedup_threshold,

model_name=self.config.embedding_model,

)

else:

accepted, duplicates = await asyncio.to_thread(

deduper.deduplicate_minhash,

accepted,

threshold=self.config.dedup_threshold,

)

for item in duplicates:

self._reject(item, "dedup", str(item.metadata.get("dedup_reason", "duplicate")))

rejected.extend(duplicates)

metrics["stages"]["dedup"] = {"method": method, "rejected": len(duplicates)}

self._observability.record_stage_outcome(

"dedup", accepted=len(accepted), rejected=len(duplicates)

)

completed.append("dedup")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "score_refine" not in completed:

with self._observability.stage_span("score_refine"):

effective_quality_mode = "off" if self._is_strict_labeled(accepted) else self.config.quality_mode

if effective_quality_mode == "off":

metrics["stages"]["score_refine"] = {

"mode": "off",

"requested_mode": self.config.quality_mode,

"scored": 0,

"refined": 0,

}

elif effective_quality_mode == "sampled":

audit_items = [

item.model_copy(deep=True)

for item in representative_sample(accepted, self.config.quality_sample_size)

]

scorer = QualityScorer(client=self.client, limiter=self.limiter, model=self.config.fast_model)

await scorer.audit_sample_batch_async(

audit_items,

batch_size=self.config.quality_audit_batch_size,

task_type=profile.task_type.value if profile and profile.task_type else None,

classes=profile.classes if profile else None,

)

scores = [

item.quality_score.overall_score

for item in audit_items

if item.quality_score is not None and not item.metadata.get("scoring_error")

]

metrics["stages"]["score_refine"] = {

"mode": "sampled",

"sample_requested": min(self.config.quality_sample_size, len(accepted)),

"scored": len(scores),

"failed": len(audit_items) - len(scores),

"remote_batches": scorer.last_audit_metrics["remote_batches"],

"remote_records": scorer.last_audit_metrics["remote_records"],

"average_overall_score": sum(scores) / len(scores) if scores else None,

"minimum_overall_score": min(scores) if scores else None,

"refined": 0,

}

else:

scorer = QualityScorer(client=self.client, limiter=self.limiter, model=self.config.fast_model)

accepted = await scorer.audit_sample_batch_async(

accepted,

batch_size=self.config.quality_audit_batch_size,

task_type=profile.task_type.value if profile and profile.task_type else None,

classes=profile.classes if profile else None,

)

refinable = [

item

for item in accepted

if item.quality_score

and item.quality_score.overall_score < self.config.refine_below_score

and item.format.value in {"alpaca", "chat", "dpo"}

and not item.metadata.get("scoring_error")

]

if refinable:

refiner = DataRefiner(client=self.client, limiter=self.limiter, model=self.model)

await refiner.refine_batch_async(refinable, mode=self.config.refine_mode)

await scorer.audit_sample_batch_async(

refinable,

batch_size=self.config.quality_audit_batch_size,

task_type=profile.task_type.value if profile and profile.task_type else None,

classes=profile.classes if profile else None,

)

metrics["stages"]["score_refine"] = {

"mode": "llm",

"scored": len(accepted),

"refined": len(refinable),

}

self._observability.record_stage_outcome("score_refine", accepted=len(accepted))

completed.append("score_refine")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "filter" not in completed:

with self._observability.stage_span("filter"):

kept: List[DatasetItem] = []

filtered: List[DatasetItem] = []

effective_quality_mode = "off" if self._is_strict_labeled(accepted) else self.config.quality_mode

if effective_quality_mode != "llm":

metrics["stages"]["filter"] = {

"action": "skipped",

"reason": (

"accuracy_contract=strict preserves labeled rows"

if self._is_strict_labeled(accepted)

else f"quality_mode={effective_quality_mode} does not produce per-row scores"

),

"accepted": len(accepted),

"rejected": 0,

}

else:

for item in accepted:

score = item.quality_score

if item.metadata.get("scoring_error"):

reason = f"Provider scoring failed: {item.metadata['scoring_error']}"

elif score is None:

reason = "No final quality score was produced."

elif not score.is_safe:

reason = "Record failed the safety assessment."

elif score.overall_score < self.config.filter_min_score:

reason = f"Final quality score {score.overall_score:.2f} is below {self.config.filter_min_score:.2f}."

else:

kept.append(item)

continue

self._reject(item, "filter", reason)

filtered.append(item)

accepted = kept

rejected.extend(filtered)

metrics["stages"]["filter"] = {"accepted": len(accepted), "rejected": len(filtered)}

self._observability.record_stage_outcome(

"filter", accepted=len(accepted), rejected=len(filtered)

)

completed.append("filter")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

if "classify" not in completed:

rejected_before = len(rejected)

with self._observability.stage_span("classify"):

classifier = DatasetClassifier(self.client, self.limiter, self.config.fast_model)

already_labeled = bool(accepted) and all(item.labels is not None for item in accepted)

schema: Optional[ClassificationSchema] = None

if self._is_strict_labeled(accepted):

action = "strict_existing_labels"

elif self.config.classification == ClassificationMode.OFF:

action = "disabled"

elif already_labeled and self.config.classification == ClassificationMode.AUTO:

action = "existing_labels"

else:

if self.config.classification == ClassificationMode.AUTO:

schema = ClassificationSchema(

applicable=bool(profile and profile.classification_applicable),

task_type=profile.task_type if profile else None,

classes=profile.classes if profile else [],

confidence=profile.confidence if profile else 0.0,

reasoning=profile.reasoning if profile else "No profile available.",

)

else:

schema = await classifier.resolve_schema(

accepted,

mode=self.config.classification,

classes=self.config.classes,

sample_size=self.config.classification_sample_size,

)

if not schema.applicable:

action = "not_applicable"

elif schema.confidence < self.config.classification_confidence:

action = "low_confidence_skip"

else:

accepted = await classifier.classify_batch(accepted, schema)

accepted, failed = classifier.partition_failures(accepted)

for item in failed:

self._reject(item, "classify", item.metadata["classification_error"])

rejected.extend(failed)

action = "classified"

metrics["stages"]["classify"] = {

"action": action,

"task_type": schema.task_type.value if schema and schema.task_type else None,

"classes": schema.classes if schema else (profile.classes if profile else []),

"confidence": schema.confidence if schema else (profile.confidence if profile else 0.0),

}

self._observability.record_stage_outcome(

"classify", accepted=len(accepted), rejected=len(rejected) - rejected_before

)

completed.append("classify")

await self._save_checkpoint(checkpoint_path, fingerprint, completed, accepted, rejected, profile, metrics)

metrics["accepted_records"] = len(accepted)

metrics["rejected_records"] = len(rejected)

metrics["accuracy_contract"] = self.config.accuracy_contract

metrics["network_policy"] = self.config.network_policy

metrics["provider"] = self.client.provider.value

metrics["model"] = self.model

metrics["usage"] = dict(getattr(self.client, "usage", {}))

usage = metrics["usage"]

if usage.get("input_tokens") or usage.get("output_tokens"):

self._observability.record_tokens(

self.client.provider.value,

self.model,

input_tokens=usage.get("input_tokens", 0),

output_tokens=usage.get("output_tokens", 0),

)

if profile is None:

profile = await DatasetProfiler(self.client, self.model).profile(accepted)

# Parallel I/O: write output, rejected, and report concurrently

write_tasks = []

if output_path:

write_tasks.append(asyncio.to_thread(write_dataset_atomic, accepted, output_path))

if rejected_path:

write_tasks.append(asyncio.to_thread(write_dataset_atomic, rejected, rejected_path))

if write_tasks:

await asyncio.gather(*write_tasks)

if report_path:

self._atomic_json(

report_path,

{

"profile": profile.model_dump(mode="json"),

"metrics": metrics,

"rejection_reasons": self._rejection_counts(rejected),

},

)

if self.config.report_html and output_path:

ReportGenerator.generate_html_report(

accepted,

output_path.with_name(f"{output_path.stem}.report.html"),

provider=f"{self.client.provider.value} / {self.model}",

)

if output_path and self._observability.prometheus_enabled:

metrics_path = output_path.with_name(f"{output_path.stem}.metrics.prom")

metrics_path.write_text(self._observability.export_prometheus_text(), encoding="utf-8")

if checkpoint_path and checkpoint_path.exists():

checkpoint_path.unlink()

return OptimizationRunResult(

accepted=accepted,

rejected=rejected,

profile=profile,

metrics=metrics,

output_path=str(output_path) if output_path else None,

rejected_path=str(rejected_path) if rejected_path else None,

report_path=str(report_path) if report_path else None,

)

@staticmethod

def _rejection_counts(items: List[DatasetItem]) -> dict[str, int]:

counts: dict[str, int] = {}

for item in items:

rejection = item.metadata.get("rejection", {})

key = f"{rejection.get('stage', 'unknown')}: {rejection.get('reason', 'unknown')}"

counts[key] = counts.get(key, 0) + 1

return counts