markitect-main/markitect/prompts/traceability/models.py

"""
Traceability models for provenance tracking.

Implements FR-11.1: Trace artifacts to producing runs, inputs, and validation.
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional


@dataclass
class RunSummary:
    """Summary of a PromptRun for traceability output."""

    run_id: str
    template_id: str
    status: str
    stage: str
    parent_run_id: Optional[str]
    depth: int
    input_bundle_hash: str
    started_at: datetime
    completed_at: Optional[datetime]

    @classmethod
    def create(
        cls,
        run_id: str,
        template_id: str,
        status: str,
        stage: str,
        input_bundle_hash: str,
        started_at: datetime,
        parent_run_id: Optional[str] = None,
        depth: int = 0,
        completed_at: Optional[datetime] = None,
    ) -> "RunSummary":
        """Create a RunSummary."""
        return cls(
            run_id=run_id,
            template_id=template_id,
            status=status,
            stage=stage,
            parent_run_id=parent_run_id,
            depth=depth,
            input_bundle_hash=input_bundle_hash,
            started_at=started_at,
            completed_at=completed_at,
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "run_id": self.run_id,
            "template_id": self.template_id,
            "status": self.status,
            "stage": self.stage,
            "parent_run_id": self.parent_run_id,
            "depth": self.depth,
            "input_bundle_hash": self.input_bundle_hash,
            "started_at": self.started_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
        }


@dataclass
class ArtifactLineage:
    """Lineage record for a single artifact in a trace."""

    artifact_id: str
    name: str
    space_id: str
    artifact_type: str
    content_digest: str
    role: str  # "input", "output", "template"

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "artifact_id": self.artifact_id,
            "name": self.name,
            "space_id": self.space_id,
            "artifact_type": self.artifact_type,
            "content_digest": self.content_digest,
            "role": self.role,
        }


@dataclass
class ProvenanceTrace:
    """Complete provenance trace for an artifact."""

    artifact_id: str
    producing_run: Optional[RunSummary] = None
    template: Optional[ArtifactLineage] = None
    input_artifacts: List[ArtifactLineage] = field(default_factory=list)
    output_artifacts: List[ArtifactLineage] = field(default_factory=list)
    generator_runs: List[RunSummary] = field(default_factory=list)
    validation_results: List[Dict[str, Any]] = field(default_factory=list)
    impact_debt: List[Dict[str, Any]] = field(default_factory=list)
    dependency_chain: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "artifact_id": self.artifact_id,
            "producing_run": self.producing_run.to_dict() if self.producing_run else None,
            "template": self.template.to_dict() if self.template else None,
            "input_artifacts": [a.to_dict() for a in self.input_artifacts],
            "output_artifacts": [a.to_dict() for a in self.output_artifacts],
            "generator_runs": [r.to_dict() for r in self.generator_runs],
            "validation_results": self.validation_results,
            "impact_debt": self.impact_debt,
            "dependency_chain": self.dependency_chain,
        }