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

osmosis_ai/cli.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Optional
+
+from .cli_commands import EvalCommand, PreviewCommand
+from .cli_services import CLIError
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    """Entry point for the osmosis CLI."""
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    handler = getattr(args, "handler", None)
+    if handler is None:
+        parser.print_help()
+        return 1
+
+    try:
+        return handler(args)
+    except CLIError as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        return 1
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="osmosis", description="Utilities for inspecting Osmosis rubric and test data files."
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    preview_parser = subparsers.add_parser(
+        "preview",
+        help="Preview a rubric YAML file or test JSONL file and print its parsed contents.",
+    )
+    PreviewCommand().configure_parser(preview_parser)
+
+    eval_parser = subparsers.add_parser(
+        "eval",
+        help="Evaluate JSONL conversations against a rubric using remote providers.",
+    )
+    EvalCommand().configure_parser(eval_parser)
+
+    return parser
+
+
+if __name__ == "__main__":
+    sys.exit(main())

osmosis_ai/cli_commands.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+from .cli_services import (
+    CLIError,
+    ParsedItem,
+    BaselineComparator,
+    ConsoleReportRenderer,
+    DatasetLoader,
+    EvaluationSession,
+    EvaluationSessionRequest,
+    JsonReportWriter,
+    RubricEvaluationEngine,
+    RubricSuite,
+    discover_rubric_config_path,
+    load_jsonl_records,
+    load_rubric_configs,
+    load_rubric_suite,
+    render_json_records,
+    render_yaml_items,
+)
+
+
+class PreviewCommand:
+    """Handler for `osmosis preview`."""
+
+    def __init__(
+        self,
+        *,
+        yaml_loader: Callable[[Path], list[ParsedItem]] = load_rubric_configs,
+        json_loader: Callable[[Path], list[dict[str, Any]]] = load_jsonl_records,
+    ):
+        self._yaml_loader = yaml_loader
+        self._json_loader = json_loader
+
+    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
+        parser.set_defaults(handler=self.run)
+        parser.add_argument(
+            "-p",
+            "--path",
+            dest="path",
+            required=True,
+            help="Path to the YAML or JSONL file to inspect.",
+        )
+
+    def run(self, args: argparse.Namespace) -> int:
+        path = Path(args.path).expanduser()
+        if not path.exists():
+            raise CLIError(f"Path '{path}' does not exist.")
+        if path.is_dir():
+            raise CLIError(f"Expected a file path but got directory '{path}'.")
+
+        suffix = path.suffix.lower()
+        if suffix in {".yaml", ".yml"}:
+            items = self._yaml_loader(path)
+            print(f"Loaded {len(items)} rubric config(s) from {path}")
+            print(render_yaml_items(items, label="Rubric config"))
+        elif suffix == ".jsonl":
+            records = self._json_loader(path)
+            print(f"Loaded {len(records)} JSONL record(s) from {path}")
+            print(render_json_records(records))
+        else:
+            raise CLIError(f"Unsupported file extension '{suffix}'. Expected .yaml, .yml, or .jsonl.")
+
+        return 0
+
+
+class EvalCommand:
+    """Handler for `osmosis eval`."""
+
+    def __init__(
+        self,
+        *,
+        session: Optional[EvaluationSession] = None,
+        config_locator: Callable[[Optional[str], Path], Path] = discover_rubric_config_path,
+        suite_loader: Callable[[Path], RubricSuite] = load_rubric_suite,
+        dataset_loader: Optional[DatasetLoader] = None,
+        engine: Optional[RubricEvaluationEngine] = None,
+        renderer: Optional[ConsoleReportRenderer] = None,
+        report_writer: Optional[JsonReportWriter] = None,
+        baseline_comparator: Optional[BaselineComparator] = None,
+    ):
+        self._renderer = renderer or ConsoleReportRenderer()
+        if session is not None:
+            self._session = session
+        else:
+            self._session = EvaluationSession(
+                config_locator=config_locator,
+                suite_loader=suite_loader,
+                dataset_loader=dataset_loader,
+                engine=engine,
+                baseline_comparator=baseline_comparator,
+                report_writer=report_writer,
+                identifier_factory=self._generate_output_identifier,
+            )
+
+    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
+        parser.set_defaults(handler=self.run)
+        parser.add_argument(
+            "-r",
+            "--rubric",
+            dest="rubric_id",
+            required=True,
+            help="Rubric identifier declared in the rubric config file.",
+        )
+        parser.add_argument(
+            "-d",
+            "--data",
+            dest="data_path",
+            required=True,
+            help="Path to the JSONL file containing evaluation records.",
+        )
+        parser.add_argument(
+            "-n",
+            "--number",
+            dest="number",
+            type=int,
+            default=1,
+            help="Run the evaluation multiple times to sample provider variance (default: 1).",
+        )
+        parser.add_argument(
+            "-c",
+            "--config",
+            dest="config_path",
+            help="Path to the rubric config YAML (defaults to searching near the data file).",
+        )
+        parser.add_argument(
+            "-o",
+            "--output",
+            dest="output_path",
+            help="Optional path to write evaluation results as JSON.",
+        )
+        parser.add_argument(
+            "-b",
+            "--baseline",
+            dest="baseline_path",
+            help="Optional path to a prior evaluation JSON to compare against.",
+        )
+
+    def run(self, args: argparse.Namespace) -> int:
+        rubric_id_raw = getattr(args, "rubric_id", "")
+        rubric_id = str(rubric_id_raw).strip()
+        if not rubric_id:
+            raise CLIError("Rubric identifier cannot be empty.")
+
+        data_path = Path(args.data_path).expanduser()
+        config_path_value = getattr(args, "config_path", None)
+        output_path_value = getattr(args, "output_path", None)
+        baseline_path_value = getattr(args, "baseline_path", None)
+
+        number_value = getattr(args, "number", None)
+        number = int(number_value) if number_value is not None else 1
+
+        request = EvaluationSessionRequest(
+            rubric_id=rubric_id,
+            data_path=data_path,
+            number=number,
+            config_path=Path(config_path_value).expanduser() if config_path_value else None,
+            output_path=Path(output_path_value).expanduser() if output_path_value else None,
+            baseline_path=Path(baseline_path_value).expanduser() if baseline_path_value else None,
+        )
+
+        try:
+            result = self._session.execute(request)
+        except KeyboardInterrupt:
+            print("Evaluation cancelled by user.")
+            return 1
+        self._renderer.render(result.report, result.baseline)
+
+        if result.written_path is not None:
+            print(f"Wrote evaluation results to {result.written_path}")
+
+        return 0
+
+    @staticmethod
+    def _generate_output_identifier() -> str:
+        return str(int(time.time()))

osmosis_ai/cli_services/__init__.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from .config import (
+    ParsedItem,
+    RubricConfig,
+    RubricConfigParser,
+    RubricSuite,
+    discover_rubric_config_path,
+    load_rubric_configs,
+    load_rubric_suite,
+    render_yaml_items,
+)
+from .dataset import (
+    ConversationMessage,
+    DatasetLoader,
+    DatasetRecord,
+    load_jsonl_records,
+    render_json_records,
+)
+from .engine import (
+    EvaluationRecordResult,
+    EvaluationReport,
+    EvaluationRun,
+    RubricEvaluationEngine,
+    RubricEvaluator,
+)
+from .errors import CLIError
+from .reporting import (
+    BaselineComparator,
+    BaselineStatistics,
+    ConsoleReportRenderer,
+    JsonReportFormatter,
+    JsonReportWriter,
+    TextReportFormatter,
+)
+from .session import EvaluationSession, EvaluationSessionRequest, EvaluationSessionResult
+
+__all__ = [
+    "BaselineComparator",
+    "BaselineStatistics",
+    "CLIError",
+    "ConsoleReportRenderer",
+    "ConversationMessage",
+    "DatasetLoader",
+    "DatasetRecord",
+    "EvaluationSession",
+    "EvaluationSessionRequest",
+    "EvaluationSessionResult",
+    "EvaluationRecordResult",
+    "EvaluationReport",
+    "EvaluationRun",
+    "JsonReportFormatter",
+    "JsonReportWriter",
+    "ParsedItem",
+    "RubricConfig",
+    "RubricConfigParser",
+    "RubricEvaluationEngine",
+    "RubricEvaluator",
+    "RubricSuite",
+    "TextReportFormatter",
+    "discover_rubric_config_path",
+    "load_jsonl_records",
+    "load_rubric_configs",
+    "load_rubric_suite",
+    "render_json_records",
+    "render_yaml_items",
+]

osmosis_ai/cli_services/config.py

@@ -0,0 +1,407 @@
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional, Sequence
+
+import yaml
+from yaml.representer import SafeRepresenter
+
+from .errors import CLIError
+from .shared import coerce_optional_float
+
+
+@dataclass(frozen=True)
+class ParsedItem:
+    label: Optional[str]
+    payload: Any
+
+
+@dataclass(frozen=True)
+class RubricConfig:
+    rubric_id: str
+    rubric_text: str
+    model_info: dict[str, Any]
+    score_min: Optional[float]
+    score_max: Optional[float]
+    system_message: Optional[str]
+    extra_info: Optional[dict[str, Any]]
+    original_input: Optional[str]
+    ground_truth: Optional[str]
+    source_label: str
+
+
+@dataclass(frozen=True)
+class RubricSuite:
+    source_path: Path
+    version: Optional[int]
+    configs: dict[str, RubricConfig]
+
+    def get(self, rubric_id: str) -> RubricConfig:
+        if rubric_id not in self.configs:
+            available = ", ".join(self.available_ids()) or "none"
+            raise CLIError(
+                f"Rubric '{rubric_id}' not found in '{self.source_path}'. Available IDs: {available}"
+            )
+        return self.configs[rubric_id]
+
+    def available_ids(self) -> list[str]:
+        return sorted(self.configs)
+
+
+@dataclass(frozen=True)
+class RubricConfigDocumentResult:
+    configs: dict[str, RubricConfig]
+    items: list[ParsedItem]
+
+
+class RubricConfigDocumentSchema:
+    """Base interface for schema-specific rubric config parsing."""
+
+    version: Optional[int] = None
+
+    def parse_document(
+        self,
+        document: Any,
+        *,
+        path: Path,
+        doc_index: int,
+        strict: bool,
+    ) -> RubricConfigDocumentResult:
+        raise NotImplementedError
+
+
+class LegacyRubricConfigSchema(RubricConfigDocumentSchema):
+    """Schema handling documents without an explicit version."""
+
+    version = None
+
+    def parse_document(
+        self,
+        document: Any,
+        *,
+        path: Path,
+        doc_index: int,
+        strict: bool,
+    ) -> RubricConfigDocumentResult:
+        defaults = _extract_config_defaults(document, path, doc_index)
+        entries = _extract_rubric_items(document, context=None, doc_index=doc_index)
+        return _build_document_configs(entries, defaults, path=path, doc_index=doc_index, strict=strict)
+
+
+class Version1RubricConfigSchema(LegacyRubricConfigSchema):
+    """Schema for version 1 documents (currently aligned with legacy layout)."""
+
+    version = 1
+
+
+class RubricConfigParser:
+    """Parses rubric configuration files and produces typed suites."""
+
+    def __init__(self, *, schemas: Optional[dict[Optional[int], RubricConfigDocumentSchema]] = None):
+        self._schemas = schemas or {
+            None: LegacyRubricConfigSchema(),
+            1: Version1RubricConfigSchema(),
+        }
+        if None not in self._schemas:
+            raise ValueError("At least one default schema (key=None) must be provided.")
+
+    def parse(self, path: Path, *, strict: bool = True) -> tuple[RubricSuite, list[ParsedItem]]:
+        documents = _load_yaml_documents(path)
+        configs: dict[str, RubricConfig] = {}
+        parsed_items: list[ParsedItem] = []
+        detected_version: Optional[int] = None
+        document_indices = []
+
+        for doc_index, document in enumerate(documents):
+            if document:
+                document_indices.append(doc_index)
+            if not document:
+                continue
+
+            doc_version = self._coerce_optional_version(document, path, doc_index)
+            if doc_version is not None:
+                if detected_version is None:
+                    detected_version = doc_version
+                elif detected_version != doc_version:
+                    raise CLIError(
+                        f"Rubric config '{path}' mixes different version numbers across documents."
+                    )
+
+        schema = self._select_schema(detected_version)
+
+        for doc_index in document_indices:
+            document = documents[doc_index]
+            if not document:
+                continue
+
+            result = schema.parse_document(
+                document,
+                path=path,
+                doc_index=doc_index,
+                strict=strict,
+            )
+            parsed_items.extend(result.items)
+            for rubric_id, config in result.configs.items():
+                if rubric_id in configs:
+                    raise CLIError(f"Duplicate rubric id '{rubric_id}' detected in '{path}'.")
+                configs[rubric_id] = config
+
+        if strict and not configs:
+            raise CLIError(f"No rubric entries found in '{path}'.")
+
+        suite = RubricSuite(source_path=path, version=detected_version, configs=configs)
+        return suite, parsed_items
+
+    def _select_schema(self, version: Optional[int]) -> RubricConfigDocumentSchema:
+        if version in self._schemas:
+            return self._schemas[version]
+        if version is None:
+            return self._schemas[None]
+        raise CLIError(f"Unsupported rubric config version '{version}'.")
+
+    @staticmethod
+    def _coerce_optional_version(document: Any, path: Path, doc_index: int) -> Optional[int]:
+        if not isinstance(document, dict):
+            return None
+        version_value = document.get("version")
+        if version_value is None:
+            return None
+        if isinstance(version_value, int):
+            if version_value < 0:
+                raise CLIError(
+                    f"Version number in '{path}' document {doc_index} must be non-negative."
+                )
+            return version_value
+        raise CLIError(
+            f"Version field in '{path}' document {doc_index} must be an integer."
+        )
+
+
+def _build_document_configs(
+    entries: Sequence[ParsedItem],
+    defaults: dict[str, Any],
+    *,
+    path: Path,
+    doc_index: int,
+    strict: bool,
+) -> RubricConfigDocumentResult:
+    configs: dict[str, RubricConfig] = {}
+    parsed_items: list[ParsedItem] = []
+
+    for item in entries:
+        payload = item.payload
+        parsed_items.append(ParsedItem(label=item.label, payload=payload))
+        if not isinstance(payload, dict):
+            continue
+
+        rubric_key_raw = payload.get("id")
+        if not isinstance(rubric_key_raw, str) or not rubric_key_raw.strip():
+            if strict:
+                raise CLIError(
+                    f"Rubric entry in '{path}' (document {doc_index}) is missing a non-empty 'id'."
+                )
+            continue
+        rubric_key = rubric_key_raw.strip()
+        if rubric_key in configs:
+            raise CLIError(f"Duplicate rubric id '{rubric_key}' detected in '{path}'.")
+
+        rubric_text = payload.get("rubric")
+        if not isinstance(rubric_text, str) or not rubric_text.strip():
+            if strict:
+                raise CLIError(
+                    f"Rubric '{rubric_key}' in '{path}' must include a non-empty 'rubric' string."
+                )
+            continue
+
+        model_info = payload.get("model_info", defaults.get("model_info"))
+        if not isinstance(model_info, dict):
+            if strict:
+                raise CLIError(
+                    f"Rubric '{rubric_key}' in '{path}' must include a 'model_info' mapping."
+                )
+            continue
+
+        extra_info_value = payload.get("extra_info", defaults.get("extra_info"))
+        if extra_info_value is not None and not isinstance(extra_info_value, dict):
+            if strict:
+                raise CLIError(
+                    f"'extra_info' for rubric '{rubric_key}' in '{path}' must be a mapping."
+                )
+            continue
+
+        try:
+            score_min = coerce_optional_float(
+                payload.get("score_min", defaults.get("score_min")),
+                "score_min",
+                f"rubric '{rubric_key}' in {path}",
+            )
+            score_max = coerce_optional_float(
+                payload.get("score_max", defaults.get("score_max")),
+                "score_max",
+                f"rubric '{rubric_key}' in {path}",
+            )
+        except CLIError:
+            if strict:
+                raise
+            continue
+
+        system_message = payload.get("system_message", defaults.get("system_message"))
+        original_input = payload.get("original_input", defaults.get("original_input"))
+        ground_truth = payload.get("ground_truth", defaults.get("ground_truth"))
+
+        label = item.label or f"document[{doc_index}]"
+        source_label = f"{path}:{label}"
+
+        configs[rubric_key] = RubricConfig(
+            rubric_id=rubric_key,
+            rubric_text=rubric_text,
+            model_info=copy.deepcopy(model_info),
+            score_min=score_min,
+            score_max=score_max,
+            system_message=system_message if isinstance(system_message, str) else None,
+            extra_info=copy.deepcopy(extra_info_value) if isinstance(extra_info_value, dict) else None,
+            original_input=original_input if isinstance(original_input, str) else None,
+            ground_truth=ground_truth if isinstance(ground_truth, str) else None,
+            source_label=source_label,
+        )
+
+    return RubricConfigDocumentResult(configs=configs, items=parsed_items)
+
+
+def discover_rubric_config_path(config_arg: Optional[str], data_path: Path) -> Path:
+    if config_arg:
+        candidate = Path(config_arg).expanduser()
+        if not candidate.exists():
+            raise CLIError(f"Rubric config path '{candidate}' does not exist.")
+        if candidate.is_dir():
+            raise CLIError(f"Rubric config path '{candidate}' is a directory.")
+        return candidate
+
+    candidates: list[Path] = []
+    candidates.append(data_path.parent / "rubric_configs.yaml")
+    candidates.append(Path.cwd() / "rubric_configs.yaml")
+    candidates.append(Path.cwd() / "examples" / "rubric_configs.yaml")
+
+    checked: list[Path] = []
+    for candidate in dict.fromkeys(candidates):
+        checked.append(candidate)
+        if candidate.exists() and candidate.is_file():
+            return candidate
+
+    searched = ", ".join(str(path) for path in checked)
+    raise CLIError(
+        "Unable to locate a rubric config file. Provide --config explicitly. "
+        f"Paths checked: {searched}"
+    )
+
+
+def load_rubric_configs(path: Path) -> list[ParsedItem]:
+    parser = RubricConfigParser()
+    _, items = parser.parse(path, strict=False)
+    return items
+
+
+def load_rubric_suite(path: Path) -> RubricSuite:
+    parser = RubricConfigParser()
+    suite, _ = parser.parse(path)
+    return suite
+
+
+def render_yaml_items(items: Sequence[ParsedItem], label: str) -> str:
+    blocks: list[str] = []
+    total = len(items)
+
+    for index, item in enumerate(items, start=1):
+        header = f"{label} #{index}"
+        if item.label:
+            header += f" ({item.label})"
+        dumped = yaml.dump(
+            item.payload,
+            Dumper=_LiteralSafeDumper,
+            sort_keys=False,
+            indent=2,
+            allow_unicode=True,
+        ).rstrip()
+
+        snippet = [header, dumped]
+        if index != total:
+            snippet.append("")
+        blocks.append("\n".join(snippet))
+
+    return "\n".join(blocks)
+
+
+def _load_yaml_documents(path: Path) -> list[Any]:
+    try:
+        with path.open("r", encoding="utf-8") as fh:
+            return list(yaml.safe_load_all(fh))
+    except yaml.YAMLError as exc:
+        raise CLIError(f"Failed to parse YAML in '{path}': {exc}") from exc
+    except OSError as exc:
+        raise CLIError(f"Unable to read rubric config '{path}': {exc}") from exc
+
+
+def _extract_config_defaults(document: Any, path: Path, doc_index: int) -> dict[str, Any]:
+    if not isinstance(document, dict):
+        return {
+            "model_info": None,
+            "extra_info": None,
+            "score_min": None,
+            "score_max": None,
+            "system_message": None,
+            "original_input": None,
+            "ground_truth": None,
+        }
+
+    source = f"document[{doc_index}] in {path}"
+
+    defaults: dict[str, Any] = {}
+    defaults["model_info"] = document.get("default_model_info")
+    defaults["extra_info"] = document.get("default_extra_info")
+    defaults["score_min"] = coerce_optional_float(
+        document.get("default_score_min"), "default_score_min", source
+    )
+    defaults["score_max"] = coerce_optional_float(
+        document.get("default_score_max"), "default_score_max", source
+    )
+    defaults["system_message"] = document.get("default_system_message")
+    defaults["original_input"] = document.get("default_original_input")
+    defaults["ground_truth"] = document.get("default_ground_truth")
+    return defaults
+
+
+def _extract_rubric_items(node: Any, context: Optional[str], doc_index: int) -> list[ParsedItem]:
+    items: list[ParsedItem] = []
+
+    if node is None:
+        return items
+
+    if isinstance(node, dict):
+        if "rubric" in node and isinstance(node["rubric"], str):
+            label = context or f"document[{doc_index}]"
+            items.append(ParsedItem(label=label, payload=node))
+        else:
+            for key, value in node.items():
+                next_context = str(key) if isinstance(key, str) else context
+                items.extend(_extract_rubric_items(value, context=next_context, doc_index=doc_index))
+    elif isinstance(node, list):
+        for index, value in enumerate(node):
+            idx_context = f"{context}[{index}]" if context else None
+            items.extend(_extract_rubric_items(value, context=idx_context, doc_index=doc_index))
+
+    return items
+
+
+class _LiteralSafeDumper(yaml.SafeDumper):
+    """YAML dumper that preserves multiline strings with literal blocks."""
+
+
+def _represent_str(dumper: yaml.Dumper, data: str):
+    if "\n" in data:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
+    return SafeRepresenter.represent_str(dumper, data)
+
+
+_LiteralSafeDumper.add_representer(str, _represent_str)