orchestrata_ultimate.core.artifact_store

Reference › orchestrata_ultimate › core › artifact_store

The canonical store for the artifacts that pipeline stages produce, shared by the kernel and the orchestrator. It closes phase 5.1 of SPEC-107 by merging two divergent stores, kernel.pipeline_artifacts and orchestrator.artifact_store, which now build on it. Kernel callers use stage and type; orchestrator callers use stage_alias and artifact_type (an enum or a string). Both spellings are accepted on construction and in queries, and internally the values are plain strings.

PipelineArtifact is a dataclass holding stage, type, content, format, optional file path, timestamp, thread and session ids, spec and requirement ids. It provides filename, to_dict (which writes both naming conventions) and from_dict. ArtifactStore keeps a lock-guarded list. add appends and, when an output_dir is configured, _persist appends the artifact to a JSONL file; on construction _replay reloads earlier artifacts from it, so old orchestrator files and new kernel files are interchangeable. Queries are by_stage, by_type, by_thread, by_requirement and by_spec. render_artifact_digest ranks artifacts by recency and optional query match and renders a digest within an artifact and token budget, using the word token counter from memory.token_counter. write_to_disk flushes artifacts into per-thread directories with extensions chosen by format.

At a glance

Class diagram

Dependency diagram

API

Classes

PipelineArtifact (dataclass)

Decorators: dataclass, traces(spec='001', story='US-006', reqs=['REQ-001-US-006-018']) · Line 69

A single artifact produced by a pipeline stage.

ArtifactStore

Decorators: traces(spec='001', story='US-006', reqs=['REQ-001-US-006-018']) · Line 196

Thread-safe collection of pipeline artifacts, queryable + persistent.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/core/artifact_store.py

"""Canonical artifact store shared by the kernel and orchestrator pipelines.

This module closes Phase 5.1 of SPEC-107: it merges the previously divergent
``kernel.pipeline_artifacts.ArtifactStore`` and
``orchestrator.artifact_store.ArtifactStore`` into a single canonical
implementation while preserving both calling conventions:

* Kernel callers use ``stage`` / ``type`` (plain strings).
* Orchestrator callers use ``stage_alias`` / ``artifact_type`` (``ArtifactType``
  enum or plain string).

Both sets of names are accepted at construction and query time; the internal
canonical representation is always ``stage`` and ``type`` strings. The JSONL
persistence format writes both key sets so older orchestrator JSONL files and
newer kernel JSONL files can be replayed interchangeably.
"""

from __future__ import annotations

import json
import logging
import os
import re
import threading
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.memory.token_counter import WordTokenCounter

_logger = logging.getLogger(__name__)

# Filename extension per artifact ``format`` so disk output is browsable.
_FORMAT_EXT = {
    "python": "py",
    "json": "json",
    "mermaid": "mmd",
    "prolog": "pl",
    "markdown": "md",
    "text": "txt",
}


def _coerce_str_type(value: Any) -> str:
    """Normalize an enum or string value to a plain string."""
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    enum_value = getattr(value, "value", None)
    if isinstance(enum_value, str):
        return enum_value
    return str(value)


def _coerce_reqs(reqs: Any) -> tuple[str, ...]:
    """Normalize requirement ids to a tuple of strings."""
    if reqs is None:
        return ()
    if isinstance(reqs, str):
        return (reqs,)
    return tuple(str(r) for r in reqs)


@dataclass
@traces(spec="001", story="US-006", reqs=["REQ-001-US-006-018"])
class PipelineArtifact:
    """A single artifact produced by a pipeline stage.

    This is the canonical value object. It accepts either kernel-style
    (``stage`` / ``type``) or orchestrator-style (``stage_alias`` /
    ``artifact_type``) field names, but stores the data as plain strings
    internally.
    """

    stage: str
    type: str
    content: str
    format: str = "text"
    file_path: str | None = None
    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
    thread_id: str = ""
    session_id: str = ""
    spec: str = ""
    reqs: tuple[str, ...] = ()

    def __init__(
        self,
        stage: str | None = None,
        stage_alias: str | None = None,
        type: str | Any | None = None,
        artifact_type: str | Any | None = None,
        content: str = "",
        format: str = "text",
        file_path: str | None = None,
        timestamp: datetime | None = None,
        thread_id: str = "",
        session_id: str = "",
        spec: str = "",
        reqs: tuple[str, ...] | list[str] | str | None = None,
    ) -> None:
        """Initialize with either kernel or orchestrator field names."""
        self.stage = stage or stage_alias or ""
        resolved_type = type if type is not None else artifact_type
        self.type = _coerce_str_type(resolved_type)
        self.content = content
        self.format = format
        self.file_path = file_path
        self.timestamp = timestamp if timestamp is not None else datetime.now(UTC)
        self.thread_id = thread_id
        self.session_id = session_id
        self.spec = spec
        self.reqs = _coerce_reqs(reqs)

    @property
    def stage_alias(self) -> str:
        """Orchestrator-style alias for :attr:`stage`."""
        return self.stage

    @property
    def artifact_type(self) -> str:
        """Orchestrator-style alias for :attr:`type`."""
        return self.type

    def filename(self, index: int = 0) -> str:
        """Build a stable on-disk filename for this artifact.

        Args:
            index: Disambiguator when multiple artifacts share stage+type.

        Returns:
            ``<stage>_<type>[_<index>].<ext>`` with the stage alias sanitised
            (``U:Ar`` → ``U-Ar``).
        """
        if self.file_path:
            return os.path.basename(self.file_path)
        ext = _FORMAT_EXT.get(self.format, "txt")
        stage = self.stage.replace(":", "-").replace("/", "-") or "stage"
        suffix = f"_{index}" if index else ""
        return f"{stage}_{self.type}{suffix}.{ext}"

    def to_dict(self) -> dict[str, Any]:
        """Serialise to a JSON-friendly dict with both naming conventions."""
        return {
            "stage": self.stage,
            "stage_alias": self.stage,
            "type": self.type,
            "artifact_type": self.type,
            "content": self.content,
            "format": self.format,
            "file_path": self.file_path,
            "timestamp": self.timestamp.isoformat(),
            "thread_id": self.thread_id,
            "session_id": self.session_id,
            "spec": self.spec,
            "reqs": list(self.reqs),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PipelineArtifact:
        """Rebuild from a JSON-friendly dict.

        Accepts either kernel-style keys (``stage`` / ``type``) or
        orchestrator-style keys (``stage_alias`` / ``artifact_type``).
        """
        ts = data.get("timestamp", "")
        try:
            timestamp = datetime.fromisoformat(ts) if ts else datetime.now(UTC)
        except ValueError:
            timestamp = datetime.now(UTC)

        reqs = data.get("reqs") or []
        if isinstance(reqs, str):
            reqs = [reqs]

        stage = data.get("stage") or data.get("stage_alias") or ""
        type_value = data.get("type") or data.get("artifact_type") or ""

        return cls(
            stage=stage,
            type=_coerce_str_type(type_value),
            content=data.get("content", ""),
            format=data.get("format", "text"),
            file_path=data.get("file_path"),
            timestamp=timestamp,
            thread_id=data.get("thread_id", ""),
            session_id=data.get("session_id", ""),
            spec=data.get("spec", ""),
            reqs=tuple(str(r) for r in reqs),
        )


@traces(spec="001", story="US-006", reqs=["REQ-001-US-006-018"])
class ArtifactStore:
    """Thread-safe collection of pipeline artifacts, queryable + persistent.

    Add artifacts as stages complete; query them by stage, type, requirement,
    or thread; flush the whole collection to disk on demand; and render a
    ranked digest within a token budget.

    When ``output_dir`` is provided, every added artifact is appended to
    ``<output_dir>/artifacts.jsonl`` and the file is replayed into the
    in-memory index on construction.
    """

    def __init__(self, output_dir: Path | str | None = None) -> None:
        """Initialise an empty, lock-guarded artifact list.

        Args:
            output_dir: Optional directory for JSONL persistence. When given,
                existing ``artifacts.jsonl`` is replayed into memory.
        """
        self._artifacts: list[PipelineArtifact] = []
        self._lock = threading.RLock()
        self._output_dir: Path | None = (
            Path(output_dir) if output_dir is not None else None
        )
        self._jsonl_path: Path | None = None
        if self._output_dir is not None:
            self._output_dir.mkdir(parents=True, exist_ok=True)
            self._jsonl_path = self._output_dir / "artifacts.jsonl"
            self._replay()

    @property
    def output_dir(self) -> Path | None:
        """The configured persistence directory, or None."""
        return self._output_dir

    # ── persistence ───────────────────────────────────────────────────────

    def _replay(self) -> None:
        """Load previously persisted artifacts from JSONL into memory."""
        if self._jsonl_path is None or not self._jsonl_path.exists():
            return
        loaded = 0
        with open(self._jsonl_path, encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                    self._artifacts.append(PipelineArtifact.from_dict(data))
                    loaded += 1
                except (json.JSONDecodeError, TypeError, KeyError) as exc:
                    _logger.warning("Skipping malformed artifact JSONL line: %s", exc)
        _logger.debug("Replayed %d artifact(s) from %s", loaded, self._jsonl_path)

    def _persist(self, artifact: PipelineArtifact) -> None:
        """Append one artifact to the JSONL backing store."""
        if self._jsonl_path is None:
            return
        try:
            with open(self._jsonl_path, "a", encoding="utf-8") as fh:
                fh.write(json.dumps(artifact.to_dict(), ensure_ascii=False))
                fh.write("\n")
        except OSError as exc:
            _logger.warning(
                "Artifact persistence failed (%s): %s", self._jsonl_path, exc
            )

    # ── mutation ──────────────────────────────────────────────────────────

    @traces(spec="001", story="US-006", reqs=["REQ-001-US-006-018"])
    def add(self, artifact: PipelineArtifact) -> None:
        """Add an artifact to the store and persist it when output_dir is set."""
        with self._lock:
            self._artifacts.append(artifact)
        self._persist(artifact)
        _logger.debug(
            "[Artifact] %s/%s (%s) thread=%s",
            artifact.stage,
            artifact.type,
            artifact.format,
            artifact.thread_id,
        )

    @traces(spec="001", story="US-006", reqs=["REQ-001-US-006-018"])
    def add_content(
        self,
        *,
        stage: str | None = None,
        stage_alias: str | None = None,
        type: str | Any | None = None,
        artifact_type: str | Any | None = None,
        content: str,
        format: str = "text",
        file_path: str | None = None,
        thread_id: str = "",
        session_id: str = "",
        spec: str = "",
        reqs: tuple[str, ...] | list[str] | None = None,
    ) -> PipelineArtifact:
        """Construct and add an artifact from its fields.

        Accepts either kernel-style (``stage`` / ``type``) or orchestrator-style
        (``stage_alias`` / ``artifact_type``) field names.

        Returns:
            The artifact that was added.
        """
        resolved_stage = stage or stage_alias or ""
        resolved_type = type if type is not None else artifact_type
        artifact = PipelineArtifact(
            stage=resolved_stage,
            type=_coerce_str_type(resolved_type),
            content=content,
            format=format,
            file_path=file_path,
            thread_id=thread_id,
            session_id=session_id,
            spec=spec,
            reqs=tuple(reqs) if reqs is not None else (),
        )
        self.add(artifact)
        return artifact

    # ── queries ───────────────────────────────────────────────────────────

    def all(self) -> list[PipelineArtifact]:
        """Return a snapshot copy of every artifact."""
        with self._lock:
            return list(self._artifacts)

    def by_stage(
        self, stage: str | None = None, *, stage_alias: str | None = None
    ) -> list[PipelineArtifact]:
        """Return all artifacts produced by *stage*."""
        target = stage or stage_alias or ""
        with self._lock:
            return [a for a in self._artifacts if a.stage == target]

    def by_type(
        self, type: str | Any | None = None, *, artifact_type: str | Any | None = None
    ) -> list[PipelineArtifact]:
        """Return all artifacts of a given logical *type*."""
        target = _coerce_str_type(type if type is not None else artifact_type)
        with self._lock:
            return [a for a in self._artifacts if a.type == target]

    def by_thread(self, thread_id: str) -> list[PipelineArtifact]:
        """Return all artifacts belonging to *thread_id*."""
        with self._lock:
            return [a for a in self._artifacts if a.thread_id == thread_id]

    def by_requirement(self, req_id: str) -> list[PipelineArtifact]:
        """Return all artifacts whose ``reqs`` include *req_id*."""
        with self._lock:
            return [a for a in self._artifacts if req_id in a.reqs]

    def by_spec(self, spec_id: str) -> list[PipelineArtifact]:
        """Return all artifacts whose ``spec`` equals *spec_id*."""
        with self._lock:
            return [a for a in self._artifacts if a.spec == spec_id]

    def count(self) -> int:
        """Return the number of artifacts in the store."""
        with self._lock:
            return len(self._artifacts)

    def __len__(self) -> int:
        with self._lock:
            return len(self._artifacts)

    # ── serialization preview ─────────────────────────────────────────────

    def to_json(self) -> str:
        """Return a compact JSON preview of every artifact.

        The output schema is preserved for backward compatibility:
        ``stage``, ``type``, ``content_preview``, ``timestamp``.
        """
        with self._lock:
            return json.dumps(
                [
                    {
                        "stage": a.stage,
                        "type": a.type,
                        "content_preview": a.content[:200],
                        "timestamp": a.timestamp.isoformat(),
                    }
                    for a in self._artifacts
                ],
                indent=2,
                default=str,
            )

    # ── ranked digest ─────────────────────────────────────────────────────

    @traces(spec="107", reqs=["REQ-107-032"])
    def render_artifact_digest(
        self,
        query: str | None = None,
        *,
        max_artifacts: int = 10,
        max_tokens: int = 2000,
    ) -> str:
        """Render a ranked markdown digest of the most relevant artifacts.

        Ranking blends recency and, when ``query`` is supplied, keyword overlap
        with the artifact's stage/type/content/spec/requirements. The digest
        stops once ``max_tokens`` is reached.

        Args:
            query: Optional free-text query to bias ranking.
            max_artifacts: Maximum artifacts to include.
            max_tokens: Token budget for the rendered digest.

        Returns:
            Markdown digest, or a short message when the store is empty.
        """
        artifacts = self.all()
        if not artifacts:
            return "_No artifacts yet._"

        ranked = self._rank_artifacts(artifacts, query or "")
        counter = WordTokenCounter()
        lines: list[str] = ["# Artifact Digest"]
        spent = counter.count(lines[0])
        count = 0
        now = datetime.now(UTC)

        for artifact in ranked[:max_artifacts]:
            age_hours = (now - artifact.timestamp).total_seconds() / 3600.0
            age_note = (
                f"{age_hours:.1f}h ago"
                if age_hours < 48
                else f"{age_hours / 24:.1f}d ago"
            )
            reqs_note = f" [{', '.join(artifact.reqs)}]" if artifact.reqs else ""
            spec_note = f" (spec:{artifact.spec})" if artifact.spec else ""
            block = (
                f"## {artifact.stage} / {artifact.type}{spec_note}{reqs_note}\n"
                f"- thread: `{artifact.thread_id or '-'}` | {age_note}\n"
                f"```\n{artifact.content[:800]}\n```"
            )
            block_tokens = counter.count(block)
            if spent + block_tokens > max_tokens and count > 0:
                break
            lines.append(block)
            spent += block_tokens
            count += 1

        return "\n\n".join(lines)

    def _rank_artifacts(
        self, artifacts: list[PipelineArtifact], query: str
    ) -> list[PipelineArtifact]:
        """Score artifacts by recency + optional keyword relevance."""
        if not query.strip():
            return sorted(artifacts, key=lambda a: a.timestamp, reverse=True)

        terms = [t.lower() for t in re.findall(r"[a-z0-9_]+", query) if len(t) > 2]
        now = datetime.now(UTC)
        scored: list[tuple[float, PipelineArtifact]] = []
        for artifact in artifacts:
            haystack = " ".join(
                [
                    artifact.stage,
                    artifact.type,
                    artifact.content,
                    artifact.spec,
                    " ".join(artifact.reqs),
                ]
            ).lower()
            matches = sum(1 for term in terms if term in haystack)
            relevance = matches / max(len(terms), 1)
            age_days = max(0.0, (now - artifact.timestamp).total_seconds() / 86400.0)
            recency = 1.0 / (1.0 + age_days)
            score = 0.6 * recency + 0.4 * relevance
            scored.append((score, artifact))
        scored.sort(key=lambda x: (x[0], x[1].timestamp), reverse=True)
        return [a for _score, a in scored]

    # ── disk flush ────────────────────────────────────────────────────────

    @traces(spec="001", story="US-006", reqs=["REQ-001-US-006-018"])
    def write_to_disk(self, output_dir: Path | str) -> list[str]:
        """Flush every artifact to ``<output_dir>/<thread_id>/``.

        Each artifact's content is written to a file named by
        :meth:`PipelineArtifact.filename`, disambiguated by index when a
        stage produces several artifacts of the same type. Artifacts with an
        empty ``thread_id`` are grouped under ``_no_thread``.

        Args:
            output_dir: Root directory (``pipeline_output`` by convention).

        Returns:
            Absolute paths of every file written.
        """
        root = Path(output_dir)
        written: list[str] = []
        with self._lock:
            artifacts = list(self._artifacts)

        seen: dict[tuple[str, str, str], int] = {}
        for artifact in artifacts:
            thread = artifact.thread_id or "_no_thread"
            target_dir = root / thread
            target_dir.mkdir(parents=True, exist_ok=True)

            key = (thread, artifact.stage, artifact.type)
            index = seen.get(key, 0)
            seen[key] = index + 1

            path = target_dir / artifact.filename(index)
            try:
                path.write_text(artifact.content, encoding="utf-8")
                written.append(str(path))
            except OSError as exc:
                _logger.warning("Artifact write failed (%s): %s", path, exc)
        return written