BioScientist/agent_system/schemas.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any
import uuid


def utc_now_iso() -> str:
    return datetime.now(tz=timezone.utc).isoformat()


@dataclass
class FailureMode:
    step_name: str
    error_type: str
    error_message: str
    hint: str = ""


@dataclass
class ExperimentReport:
    run_id: str
    created_at: str
    task: str
    status: str
    input_manifest: dict[str, Any]
    pipeline_config_id: str | None
    output_artifacts: list[str]
    metrics: dict[str, Any]
    execution_log: str
    failures: list[FailureMode] = field(default_factory=list)
    notes: str = ""

    @staticmethod
    def new(
        task: str,
        input_manifest: dict[str, Any],
        pipeline_config_id: str | None = None,
    ) -> "ExperimentReport":
        return ExperimentReport(
            run_id=f"exp-{uuid.uuid4().hex[:12]}",
            created_at=utc_now_iso(),
            task=task,
            status="running",
            input_manifest=input_manifest,
            pipeline_config_id=pipeline_config_id,
            output_artifacts=[],
            metrics={},
            execution_log="",
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class Insight:
    insight_id: str
    created_at: str
    title: str
    hypothesis: str
    recommendation: str
    confidence: float
    evidence_run_ids: list[str]
    tags: list[str] = field(default_factory=list)

    @staticmethod
    def build(
        title: str,
        hypothesis: str,
        recommendation: str,
        confidence: float,
        evidence_run_ids: list[str],
        tags: list[str] | None = None,
    ) -> "Insight":
        return Insight(
            insight_id=f"ins-{uuid.uuid4().hex[:12]}",
            created_at=utc_now_iso(),
            title=title,
            hypothesis=hypothesis,
            recommendation=recommendation,
            confidence=max(0.0, min(1.0, confidence)),
            evidence_run_ids=evidence_run_ids,
            tags=tags or [],
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class PipelineConfiguration:
    config_id: str
    created_at: str
    strategy_name: str
    task_scope: str
    tools: list[str]
    parameters: dict[str, Any]
    rationale: str
    source_insight_ids: list[str] = field(default_factory=list)

    @staticmethod
    def build(
        strategy_name: str,
        task_scope: str,
        tools: list[str],
        parameters: dict[str, Any],
        rationale: str,
        source_insight_ids: list[str] | None = None,
    ) -> "PipelineConfiguration":
        return PipelineConfiguration(
            config_id=f"cfg-{uuid.uuid4().hex[:12]}",
            created_at=utc_now_iso(),
            strategy_name=strategy_name,
            task_scope=task_scope,
            tools=tools,
            parameters=parameters,
            rationale=rationale,
            source_insight_ids=source_insight_ids or [],
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class Hypothesis:
    hypothesis_id: str
    created_at: str
    domain: str
    user_query: str
    title: str
    hypothesis: str
    expected_improvement: str
    theoretical_basis: str
    tags: list[str] = field(default_factory=list)
    source_examples: list[str] = field(default_factory=list)
    data_sources: list[str] = field(default_factory=list)
    data_operations: list[dict[str, Any]] = field(default_factory=list)
    historical_reflection: list[str] = field(default_factory=list)
    error_avoidance: list[str] = field(default_factory=list)
    reasoning_chain: list[str] = field(default_factory=list)

    @staticmethod
    def build(
        domain: str,
        user_query: str,
        title: str,
        hypothesis: str,
        expected_improvement: str,
        theoretical_basis: str,
        tags: list[str] | None = None,
        source_examples: list[str] | None = None,
        data_sources: list[str] | None = None,
        data_operations: list[dict[str, Any]] | None = None,
        historical_reflection: list[str] | None = None,
        error_avoidance: list[str] | None = None,
        reasoning_chain: list[str] | None = None,
    ) -> "Hypothesis":
        return Hypothesis(
            hypothesis_id=f"hyp-{uuid.uuid4().hex[:12]}",
            created_at=utc_now_iso(),
            domain=domain,
            user_query=user_query,
            title=title,
            hypothesis=hypothesis,
            expected_improvement=expected_improvement,
            theoretical_basis=theoretical_basis,
            tags=tags or [],
            source_examples=source_examples or [],
            data_sources=data_sources or [],
            data_operations=data_operations or [],
            historical_reflection=historical_reflection or [],
            error_avoidance=error_avoidance or [],
            reasoning_chain=reasoning_chain or [],
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class ValidationReport:
    validation_id: str
    created_at: str
    hypothesis_id: str
    domain: str
    level: str
    status: str
    score: float
    confidence: float
    key_metrics: dict[str, Any]
    failure_reason: str
    evidence: str

    @staticmethod
    def build(
        hypothesis_id: str,
        domain: str,
        level: str,
        status: str,
        score: float,
        confidence: float,
        key_metrics: dict[str, Any] | None = None,
        failure_reason: str = "",
        evidence: str = "",
    ) -> "ValidationReport":
        normalized_status = status if status in {"success", "failed", "inconclusive"} else "inconclusive"
        return ValidationReport(
            validation_id=f"val-{uuid.uuid4().hex[:12]}",
            created_at=utc_now_iso(),
            hypothesis_id=hypothesis_id,
            domain=domain,
            level=level,
            status=normalized_status,
            score=max(0.0, min(1.0, float(score))),
            confidence=max(0.0, min(1.0, float(confidence))),
            key_metrics=key_metrics or {},
            failure_reason=failure_reason,
            evidence=evidence,
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)