osmosis-ai 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

osmosis_ai/cli_services/dataset.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import copy
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional, Sequence
+
+from .errors import CLIError
+from .shared import coerce_optional_float, gather_text_fragments
+
+
+@dataclass(frozen=True)
+class ConversationMessage:
+    """Normalized conversation message with preserved raw payload fields."""
+
+    role: str
+    content: Any
+    metadata: dict[str, Any]
+
+    def to_payload(self) -> dict[str, Any]:
+        payload: dict[str, Any] = copy.deepcopy(self.metadata)
+        payload["role"] = self.role
+        if self.content is None:
+            payload.pop("content", None)
+        else:
+            payload["content"] = copy.deepcopy(self.content)
+        return payload
+
+    def text_fragments(self) -> list[str]:
+        fragments: list[str] = []
+        seen: set[int] = set()
+        gather_text_fragments(self.content, fragments, allow_free_strings=True, seen=seen)
+        for value in self.metadata.values():
+            gather_text_fragments(value, fragments, seen=seen)
+        return fragments
+
+    @classmethod
+    def from_raw(cls, raw: dict[str, Any], *, source_label: str, index: int) -> "ConversationMessage":
+        role_value = raw.get("role")
+        if not isinstance(role_value, str) or not role_value.strip():
+            raise CLIError(
+                f"Message {index} in {source_label} must include a non-empty string 'role'."
+            )
+        content_value = copy.deepcopy(raw.get("content"))
+        metadata: dict[str, Any] = {}
+        for key, value in raw.items():
+            if key in {"role", "content"}:
+                continue
+            metadata[str(key)] = copy.deepcopy(value)
+        return cls(role=role_value.strip().lower(), content=content_value, metadata=metadata)
+
+
+@dataclass(frozen=True)
+class DatasetRecord:
+    payload: dict[str, Any]
+    rubric_id: str
+    conversation_id: Optional[str]
+    record_id: Optional[str]
+    messages: tuple[ConversationMessage, ...]
+    ground_truth: Optional[str]
+    system_message: Optional[str]
+    original_input: Optional[str]
+    metadata: Optional[dict[str, Any]]
+    extra_info: Optional[dict[str, Any]]
+    score_min: Optional[float]
+    score_max: Optional[float]
+
+    def message_payloads(self) -> list[dict[str, Any]]:
+        """Return messages as provider-ready payloads."""
+        return [message.to_payload() for message in self.messages]
+
+    def merged_extra_info(self, config_extra: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:
+        merged: dict[str, Any] = {}
+        if isinstance(config_extra, dict):
+            merged.update(copy.deepcopy(config_extra))
+        if isinstance(self.extra_info, dict):
+            merged.update(copy.deepcopy(self.extra_info))
+        if isinstance(self.metadata, dict) and self.metadata:
+            merged.setdefault("dataset_metadata", copy.deepcopy(self.metadata))
+        return merged or None
+
+    def assistant_preview(self, *, max_length: int = 140) -> Optional[str]:
+        for message in reversed(self.messages):
+            if message.role != "assistant":
+                continue
+            fragments = message.text_fragments()
+            if not fragments:
+                continue
+            preview = " ".join(" ".join(fragments).split())
+            if not preview:
+                continue
+            if len(preview) > max_length:
+                preview = preview[: max_length - 3].rstrip() + "..."
+            return preview
+        return None
+
+    def conversation_label(self, fallback_index: int) -> str:
+        if isinstance(self.conversation_id, str) and self.conversation_id.strip():
+            return self.conversation_id.strip()
+        return f"record[{fallback_index}]"
+
+    def record_identifier(self, conversation_label: str) -> str:
+        if isinstance(self.record_id, str) and self.record_id.strip():
+            return self.record_id.strip()
+        raw_id = self.payload.get("id")
+        if isinstance(raw_id, str) and raw_id.strip():
+            return raw_id.strip()
+        if raw_id is not None:
+            return str(raw_id)
+        return conversation_label
+
+
+class DatasetLoader:
+    """Loads dataset records from JSONL files."""
+
+    def load(self, path: Path) -> list[DatasetRecord]:
+        records: list[DatasetRecord] = []
+        with path.open("r", encoding="utf-8") as fh:
+            for line_number, raw_line in enumerate(fh, start=1):
+                stripped = raw_line.strip()
+                if not stripped:
+                    continue
+                try:
+                    payload = json.loads(stripped)
+                except json.JSONDecodeError as exc:
+                    raise CLIError(
+                        f"Invalid JSON on line {line_number} of '{path}': {exc.msg}"
+                    ) from exc
+                if not isinstance(payload, dict):
+                    raise CLIError(
+                        f"Expected JSON object on line {line_number} of '{path}'."
+                    )
+
+                records.append(self._create_record(payload))
+
+        if not records:
+            raise CLIError(f"No JSON records found in '{path}'.")
+
+        return records
+
+    @staticmethod
+    def _create_record(payload: dict[str, Any]) -> DatasetRecord:
+        rubric_id = payload.get("rubric_id")
+        rubric_id_str = str(rubric_id).strip() if isinstance(rubric_id, str) else ""
+
+        conversation_id_raw = payload.get("conversation_id")
+        conversation_id = None
+        if isinstance(conversation_id_raw, str) and conversation_id_raw.strip():
+            conversation_id = conversation_id_raw.strip()
+
+        record_id_raw = payload.get("id")
+        record_id = str(record_id_raw).strip() if isinstance(record_id_raw, str) else None
+
+        score_min = coerce_optional_float(
+            payload.get("score_min"), "score_min", f"record '{conversation_id or rubric_id or '<record>'}'"
+        )
+        score_max = coerce_optional_float(
+            payload.get("score_max"), "score_max", f"record '{conversation_id or rubric_id or '<record>'}'"
+        )
+
+        metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else None
+        extra_info = payload.get("extra_info") if isinstance(payload.get("extra_info"), dict) else None
+        record_label = conversation_id or record_id or rubric_id_str or "<record>"
+        messages = _parse_messages(payload.get("messages"), source_label=record_label)
+
+        return DatasetRecord(
+            payload=payload,
+            rubric_id=rubric_id_str,
+            conversation_id=conversation_id,
+            record_id=record_id,
+            messages=messages,
+            ground_truth=payload.get("ground_truth") if isinstance(payload.get("ground_truth"), str) else None,
+            system_message=payload.get("system_message") if isinstance(payload.get("system_message"), str) else None,
+            original_input=payload.get("original_input") if isinstance(payload.get("original_input"), str) else None,
+            metadata=metadata,
+            extra_info=extra_info,
+            score_min=score_min,
+            score_max=score_max,
+        )
+
+
+def load_jsonl_records(path: Path) -> list[dict[str, Any]]:
+    records: list[dict[str, Any]] = []
+    with path.open("r", encoding="utf-8") as fh:
+        for line_number, raw_line in enumerate(fh, start=1):
+            stripped = raw_line.strip()
+            if not stripped:
+                continue
+            try:
+                record = json.loads(stripped)
+            except json.JSONDecodeError as exc:
+                raise CLIError(f"Invalid JSON on line {line_number} of '{path}': {exc.msg}") from exc
+            if not isinstance(record, dict):
+                raise CLIError(f"Expected JSON object on line {line_number} of '{path}'.")
+            records.append(record)
+
+    if not records:
+        raise CLIError(f"No JSON records found in '{path}'.")
+
+    return records
+
+
+def render_json_records(records: Sequence[dict[str, Any]]) -> str:
+    segments: list[str] = []
+    total = len(records)
+
+    for index, record in enumerate(records, start=1):
+        body = json.dumps(record, indent=2, ensure_ascii=False)
+        snippet = [f"JSONL record #{index}", body]
+        if index != total:
+            snippet.append("")
+        segments.append("\n".join(snippet))
+
+    return "\n".join(segments)
+
+
+def _parse_messages(messages: Any, *, source_label: str) -> tuple[ConversationMessage, ...]:
+    if not isinstance(messages, list) or not messages:
+        raise CLIError(f"Record '{source_label}' must include a non-empty 'messages' list.")
+
+    normalized: list[ConversationMessage] = []
+    for index, entry in enumerate(messages):
+        if not isinstance(entry, dict):
+            raise CLIError(
+                f"Message {index} in {source_label} must be an object, got {type(entry).__name__}."
+            )
+        normalized.append(ConversationMessage.from_raw(entry, source_label=source_label, index=index))
+    return tuple(normalized)

osmosis_ai/cli_services/engine.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import copy
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional, Sequence
+
+from tqdm import tqdm
+
+from ..rubric_eval import evaluate_rubric
+from ..rubric_types import MissingAPIKeyError, ModelNotFoundError, ProviderRequestError
+from .config import RubricConfig
+from .dataset import DatasetRecord
+from .errors import CLIError
+from .shared import calculate_statistics, coerce_optional_float, collapse_preview_text
+
+
+class RubricEvaluator:
+    """Thin wrapper over evaluate_rubric to enable injection during tests."""
+
+    def __init__(self, evaluate_fn: Any = evaluate_rubric):
+        self._evaluate_fn = evaluate_fn
+
+    def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
+        messages = record.message_payloads()
+        if not messages:
+            label = record.conversation_id or record.rubric_id or "<record>"
+            raise CLIError(f"Record '{label}' must include a non-empty 'messages' list.")
+
+        score_min = coerce_optional_float(
+            record.score_min if record.score_min is not None else config.score_min,
+            "score_min",
+            f"record '{record.conversation_id or '<record>'}'",
+        )
+        score_max = coerce_optional_float(
+            record.score_max if record.score_max is not None else config.score_max,
+            "score_max",
+            f"record '{record.conversation_id or '<record>'}'",
+        )
+
+        try:
+            return self._evaluate_fn(
+                rubric=config.rubric_text,
+                messages=messages,
+                model_info=copy.deepcopy(config.model_info),
+                ground_truth=record.ground_truth if record.ground_truth is not None else config.ground_truth,
+                system_message=record.system_message if record.system_message is not None else config.system_message,
+                original_input=record.original_input if record.original_input is not None else config.original_input,
+                extra_info=record.merged_extra_info(config.extra_info),
+                score_min=score_min,
+                score_max=score_max,
+                return_details=True,
+            )
+        except (MissingAPIKeyError, ProviderRequestError, ModelNotFoundError) as exc:
+            raise CLIError(str(exc)) from exc
+
+
+@dataclass
+class EvaluationRun:
+    run_index: int
+    status: str
+    score: Optional[float]
+    explanation: Optional[str]
+    preview: Optional[str]
+    duration_seconds: float
+    started_at: datetime
+    completed_at: datetime
+    error: Optional[str]
+    raw: Any
+
+
+@dataclass
+class EvaluationRecordResult:
+    record_index: int
+    record: DatasetRecord
+    conversation_label: str
+    runs: list[EvaluationRun]
+    statistics: dict[str, float]
+
+
+@dataclass
+class EvaluationReport:
+    rubric_config: RubricConfig
+    config_path: Path
+    data_path: Path
+    number: int
+    record_results: list[EvaluationRecordResult]
+    overall_statistics: dict[str, float]
+
+
+class RubricEvaluationEngine:
+    """Executes rubric evaluations across a dataset and aggregates statistics."""
+
+    def __init__(self, evaluator: Optional[RubricEvaluator] = None):
+        self._evaluator = evaluator or RubricEvaluator()
+
+    def execute(
+        self,
+        *,
+        rubric_config: RubricConfig,
+        config_path: Path,
+        data_path: Path,
+        records: Sequence[DatasetRecord],
+        number: int,
+    ) -> EvaluationReport:
+        record_results: list[EvaluationRecordResult] = []
+        aggregate_scores: list[float] = []
+        total_runs = 0
+        total_successes = 0
+
+        progress_total = len(records) * number
+        show_progress = progress_total > 1 and getattr(sys.stderr, "isatty", lambda: False)()
+        progress = (
+            tqdm(
+                total=progress_total,
+                file=sys.stderr,
+                dynamic_ncols=True,
+                leave=False,
+            )
+            if show_progress
+            else None
+        )
+
+        try:
+            for record_index, record in enumerate(records, start=1):
+                conversation_label = record.conversation_label(record_index)
+                fallback_preview = record.assistant_preview()
+
+                runs: list[EvaluationRun] = []
+                scores: list[float] = []
+
+                for attempt in range(1, number + 1):
+                    started_at = datetime.now(timezone.utc)
+                    timer_start = time.perf_counter()
+                    status = "success"
+                    error_message: Optional[str] = None
+                    score_value: Optional[float] = None
+                    explanation_value: Optional[str] = None
+                    preview_value: Optional[str] = None
+                    raw_payload: Any = None
+
+                    try:
+                        result = self._evaluator.run(rubric_config, record)
+                    except CLIError as exc:
+                        status = "error"
+                        error_message = str(exc)
+                        result = None
+                    except Exception as exc:  # pragma: no cover - unexpected path
+                        status = "error"
+                        error_message = f"{type(exc).__name__}: {exc}"
+                        result = None
+
+                    duration_seconds = time.perf_counter() - timer_start
+                    completed_at = datetime.now(timezone.utc)
+
+                    if status == "success" and isinstance(result, dict):
+                        raw_payload = result.get("raw")
+                        score_value = _extract_float(result.get("score"))
+                        explanation_value = _normalize_optional_text(result.get("explanation"))
+                        preview_value = self._resolve_preview_text(result, fallback_preview)
+                        if score_value is not None:
+                            scores.append(score_value)
+                            aggregate_scores.append(score_value)
+                            total_successes += 1
+                    else:
+                        preview_value = fallback_preview
+
+                    total_runs += 1
+
+                    runs.append(
+                        EvaluationRun(
+                            run_index=attempt,
+                            status=status,
+                            score=score_value,
+                            explanation=explanation_value,
+                            preview=preview_value,
+                            duration_seconds=duration_seconds,
+                            started_at=started_at,
+                            completed_at=completed_at,
+                            error=error_message,
+                            raw=raw_payload,
+                        )
+                    )
+
+                    if progress:
+                        progress.update()
+
+                statistics = calculate_statistics(scores)
+                statistics["total_runs"] = len(runs)
+                statistics["success_count"] = len(scores)
+                statistics["failure_count"] = len(runs) - len(scores)
+                record_results.append(
+                    EvaluationRecordResult(
+                        record_index=record_index,
+                        record=record,
+                        conversation_label=conversation_label,
+                        runs=runs,
+                        statistics=statistics,
+                    )
+                )
+        finally:
+            if progress:
+                progress.close()
+
+        overall_statistics = calculate_statistics(aggregate_scores)
+        overall_statistics["total_runs"] = total_runs
+        overall_statistics["success_count"] = total_successes
+        overall_statistics["failure_count"] = total_runs - total_successes
+        return EvaluationReport(
+            rubric_config=rubric_config,
+            config_path=config_path,
+            data_path=data_path,
+            number=number,
+            record_results=record_results,
+            overall_statistics=overall_statistics,
+        )
+
+    @staticmethod
+    def _resolve_preview_text(result: Optional[dict[str, Any]], fallback: Optional[str]) -> Optional[str]:
+        if not isinstance(result, dict):
+            return fallback
+        preview = collapse_preview_text(result.get("preview"))
+        if preview:
+            return preview
+
+        raw_payload = result.get("raw")
+        if isinstance(raw_payload, dict):
+            for key in ("preview", "summary", "text"):
+                preview = collapse_preview_text(raw_payload.get(key))
+                if preview:
+                    return preview
+        return fallback
+
+
+def _extract_float(value: Any) -> Optional[float]:
+    try:
+        if isinstance(value, (int, float)) and not isinstance(value, bool):
+            return float(value)
+        return None
+    except (TypeError, ValueError):
+        return None
+
+
+def _normalize_optional_text(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text or None

osmosis_ai/cli_services/errors.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+
+class CLIError(Exception):
+    """Raised when the CLI encounters a recoverable error."""
+
+    pass