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

osmosis_ai/cli_services/session.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, Sequence
+
+from ..rubric_eval import ensure_api_key_available
+from ..rubric_types import MissingAPIKeyError
+from .config import RubricConfig, RubricSuite, discover_rubric_config_path, load_rubric_suite
+from .dataset import DatasetLoader, DatasetRecord
+from .engine import RubricEvaluationEngine, EvaluationReport
+from .errors import CLIError
+from .reporting import BaselineComparator, BaselineStatistics, JsonReportWriter
+
+
+_CACHE_ROOT = Path("~/.cache/osmosis/eval_result").expanduser()
+
+
+def _sanitise_rubric_folder(rubric_id: str) -> str:
+    """Produce a filesystem-safe folder name for the rubric id."""
+    clean = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in rubric_id.strip())
+    clean = clean.strip("_") or "rubric"
+    return clean.lower()
+
+
+@dataclass(frozen=True)
+class EvaluationSessionRequest:
+    rubric_id: str
+    data_path: Path
+    number: int = 1
+    config_path: Optional[Path] = None
+    output_path: Optional[Path] = None
+    output_identifier: Optional[str] = None
+    baseline_path: Optional[Path] = None
+
+
+@dataclass
+class EvaluationSessionResult:
+    request: EvaluationSessionRequest
+    config_path: Path
+    data_path: Path
+    rubric_config: RubricConfig
+    records: Sequence[DatasetRecord]
+    report: EvaluationReport
+    baseline: Optional[BaselineStatistics]
+    written_path: Optional[Path]
+    output_identifier: Optional[str]
+
+
+class EvaluationSession:
+    """Coordinates rubric evaluation end-to-end for reusable orchestration."""
+
+    def __init__(
+        self,
+        *,
+        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,
+        baseline_comparator: Optional[BaselineComparator] = None,
+        report_writer: Optional[JsonReportWriter] = None,
+        identifier_factory: Optional[Callable[[], str]] = None,
+    ):
+        self._config_locator = config_locator
+        self._suite_loader = suite_loader
+        self._dataset_loader = dataset_loader or DatasetLoader()
+        self._engine = engine or RubricEvaluationEngine()
+        self._baseline_comparator = baseline_comparator or BaselineComparator()
+        self._report_writer = report_writer or JsonReportWriter()
+        self._identifier_factory = identifier_factory or self._default_identifier
+
+    def execute(self, request: EvaluationSessionRequest) -> EvaluationSessionResult:
+        rubric_id = request.rubric_id.strip()
+        if not rubric_id:
+            raise CLIError("Rubric identifier cannot be empty.")
+
+        number_value = request.number if request.number is not None else 1
+        number = int(number_value)
+        if number < 1:
+            raise CLIError("Number of runs must be a positive integer.")
+
+        data_path = request.data_path.expanduser()
+        if not data_path.exists():
+            raise CLIError(f"Data path '{data_path}' does not exist.")
+        if data_path.is_dir():
+            raise CLIError(f"Expected a JSONL file but received directory '{data_path}'.")
+
+        config_override = str(request.config_path.expanduser()) if request.config_path else None
+        config_path = self._config_locator(config_override, data_path)
+        suite = self._suite_loader(config_path)
+        rubric_config = suite.get(rubric_id)
+
+        try:
+            ensure_api_key_available(rubric_config.model_info)
+        except (MissingAPIKeyError, TypeError) as exc:
+            raise CLIError(str(exc)) from exc
+
+        all_records = self._dataset_loader.load(data_path)
+        matching_records = [
+            record for record in all_records if record.rubric_id.lower() == rubric_id.lower()
+        ]
+        if not matching_records:
+            raise CLIError(f"No records in '{data_path}' reference rubric '{rubric_id}'.")
+
+        baseline_stats = self._load_baseline(request.baseline_path)
+
+        resolved_output_path, resolved_identifier = self._resolve_output_path(
+            request.output_path,
+            request.output_identifier,
+            rubric_id=rubric_id,
+        )
+
+        report = self._engine.execute(
+            rubric_config=rubric_config,
+            config_path=config_path,
+            data_path=data_path,
+            records=matching_records,
+            number=number,
+        )
+
+        written_path = None
+        if resolved_output_path is not None:
+            written_path = self._report_writer.write(
+                report,
+                output_path=resolved_output_path,
+                output_identifier=resolved_identifier,
+                baseline=baseline_stats,
+            )
+
+        return EvaluationSessionResult(
+            request=request,
+            config_path=config_path,
+            data_path=data_path,
+            rubric_config=rubric_config,
+            records=matching_records,
+            report=report,
+            baseline=baseline_stats,
+            written_path=written_path,
+            output_identifier=resolved_identifier,
+        )
+
+    def _load_baseline(self, baseline_path: Optional[Path]) -> Optional[BaselineStatistics]:
+        if baseline_path is None:
+            return None
+        resolved = baseline_path.expanduser()
+        return self._baseline_comparator.load(resolved)
+
+    def _resolve_output_path(
+        self,
+        output_candidate: Optional[Path],
+        output_identifier: Optional[str],
+        *,
+        rubric_id: str,
+    ) -> tuple[Optional[Path], Optional[str]]:
+        if output_candidate is None:
+            identifier = output_identifier or self._identifier_factory()
+            target_dir = _CACHE_ROOT / _sanitise_rubric_folder(rubric_id)
+            target_dir.mkdir(parents=True, exist_ok=True)
+            return target_dir / f"rubric_eval_result_{identifier}.json", identifier
+
+        candidate = output_candidate.expanduser()
+        if candidate.suffix:
+            if candidate.exists() and candidate.is_dir():
+                raise CLIError(f"Output path '{candidate}' is a directory.")
+            return candidate, output_identifier
+
+        candidate.mkdir(parents=True, exist_ok=True)
+        identifier = output_identifier or self._identifier_factory()
+        return candidate / f"rubric_eval_result_{identifier}.json", identifier
+
+    @staticmethod
+    def _default_identifier() -> str:
+        return str(int(time.time()))

osmosis_ai/cli_services/shared.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from statistics import mean, pvariance, pstdev
+from typing import Any, Collection, Optional, Set
+
+from .errors import CLIError
+
+
+def coerce_optional_float(value: Any, field_name: str, source_label: str) -> Optional[float]:
+    if value is None:
+        return None
+    if isinstance(value, (int, float)) and not isinstance(value, bool):
+        return float(value)
+    raise CLIError(
+        f"Expected '{field_name}' in {source_label} to be numeric, got {type(value).__name__}."
+    )
+
+
+def collapse_preview_text(value: Any, *, max_length: int = 140) -> Optional[str]:
+    if not isinstance(value, str):
+        return None
+    collapsed = " ".join(value.strip().split())
+    if not collapsed:
+        return None
+    if len(collapsed) > max_length:
+        collapsed = collapsed[: max_length - 3].rstrip() + "..."
+    return collapsed
+
+
+def calculate_statistics(scores: list[float]) -> dict[str, float]:
+    if not scores:
+        return {
+            "average": 0.0,
+            "variance": 0.0,
+            "stdev": 0.0,
+            "min": 0.0,
+            "max": 0.0,
+        }
+    average = mean(scores)
+    variance = pvariance(scores)
+    std_dev = pstdev(scores)
+    return {
+        "average": average,
+        "variance": variance,
+        "stdev": std_dev,
+        "min": min(scores),
+        "max": max(scores),
+    }
+
+
+def calculate_stat_deltas(baseline: dict[str, float], current: dict[str, float]) -> dict[str, float]:
+    delta: dict[str, float] = {}
+    for key, current_value in current.items():
+        if key not in baseline:
+            continue
+        try:
+            baseline_value = float(baseline[key])
+            current_numeric = float(current_value)
+        except (TypeError, ValueError):
+            continue
+        delta[key] = current_numeric - baseline_value
+    return delta
+
+
+def gather_text_fragments(
+    node: Any,
+    fragments: list[str],
+    *,
+    allow_free_strings: bool = False,
+    seen: Optional[Set[int]] = None,
+    string_key_allowlist: Optional[Collection[str]] = None,
+) -> None:
+    """Collect textual snippets from nested message-like structures.
+
+    The traversal favours common chat-completions shapes (e.g. ``{"type": "text"}``
+    blocks) and avoids indiscriminately pulling in metadata values such as IDs.
+    ``allow_free_strings`` controls whether bare strings encountered at the current
+    level should be considered textual content (useful for raw message content but
+    typically disabled for metadata fields).
+    """
+
+    if seen is None:
+        seen = set()
+
+    if isinstance(node, str):
+        if allow_free_strings:
+            stripped = node.strip()
+            if stripped:
+                fragments.append(stripped)
+        return
+
+    if isinstance(node, list):
+        for item in node:
+            gather_text_fragments(
+                item,
+                fragments,
+                allow_free_strings=allow_free_strings,
+                seen=seen,
+                string_key_allowlist=string_key_allowlist,
+            )
+        return
+
+    if not isinstance(node, dict):
+        return
+
+    node_id = id(node)
+    if node_id in seen:
+        return
+    seen.add(node_id)
+
+    allowlist = {"text", "value", "message"}
+    if string_key_allowlist is not None:
+        allowlist = {key.lower() for key in string_key_allowlist}
+    else:
+        allowlist = {key.lower() for key in allowlist}
+
+    prioritized_keys = ("text", "value")
+    handled_keys: Set[str] = {
+        "text",
+        "value",
+        "content",
+        "message",
+        "parts",
+        "input_text",
+        "output_text",
+        "type",
+        "role",
+        "name",
+        "id",
+        "index",
+        "finish_reason",
+        "reason",
+        "tool_call_id",
+        "metadata",
+    }
+
+    for key in prioritized_keys:
+        if key not in node:
+            continue
+        before_count = len(fragments)
+        gather_text_fragments(
+            node[key],
+            fragments,
+            allow_free_strings=True,
+            seen=seen,
+            string_key_allowlist=string_key_allowlist,
+        )
+        if len(fragments) > before_count:
+            break
+
+    if node.get("type") == "tool_result" and "content" in node:
+        gather_text_fragments(
+            node["content"],
+            fragments,
+            allow_free_strings=True,
+            seen=seen,
+            string_key_allowlist=string_key_allowlist,
+        )
+    elif "content" in node:
+        gather_text_fragments(
+            node["content"],
+            fragments,
+            allow_free_strings=True,
+            seen=seen,
+            string_key_allowlist=string_key_allowlist,
+        )
+
+    for key in ("message", "parts", "input_text", "output_text"):
+        if key in node:
+            gather_text_fragments(
+                node[key],
+                fragments,
+                allow_free_strings=True,
+                seen=seen,
+                string_key_allowlist=string_key_allowlist,
+            )
+
+    for key, value in node.items():
+        if key in handled_keys:
+            continue
+        if isinstance(value, (list, dict)):
+            gather_text_fragments(
+                value,
+                fragments,
+                allow_free_strings=False,
+                seen=seen,
+                string_key_allowlist=string_key_allowlist,
+            )
+        elif isinstance(value, str) and key.lower() in allowlist:
+            stripped = value.strip()
+            if stripped:
+                fragments.append(stripped)
+
+
+def collect_text_fragments(
+    node: Any,
+    *,
+    allow_free_strings: bool = False,
+    string_key_allowlist: Optional[Collection[str]] = None,
+) -> list[str]:
+    fragments: list[str] = []
+    gather_text_fragments(
+        node,
+        fragments,
+        allow_free_strings=allow_free_strings,
+        seen=set(),
+        string_key_allowlist=string_key_allowlist,
+    )
+    return fragments

osmosis_ai/consts.py

@@ -1,3 +1,3 @@
 # package metadata
 package_name = "osmosis-ai"
-package_version = "0.2.1"
+package_version = "0.2.2"

osmosis_ai/providers/__init__.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Tuple
+
+from .anthropic_provider import AnthropicProvider
+from .base import DEFAULT_REQUEST_TIMEOUT_SECONDS, ProviderRegistry, ProviderRequest, RubricProvider
+from .gemini_provider import GeminiProvider
+from .openai_family import OpenAIProvider, XAIProvider
+
+_REGISTRY = ProviderRegistry()
+_REGISTRY.register(OpenAIProvider())
+_REGISTRY.register(XAIProvider())
+_REGISTRY.register(AnthropicProvider())
+_REGISTRY.register(GeminiProvider())
+
+
+def get_provider(name: str) -> RubricProvider:
+    return _REGISTRY.get(name)
+
+
+def register_provider(provider: RubricProvider) -> None:
+    _REGISTRY.register(provider)
+
+
+def supported_providers() -> Tuple[str, ...]:
+    return _REGISTRY.supported_providers()
+
+
+__all__ = [
+    "DEFAULT_REQUEST_TIMEOUT_SECONDS",
+    "ProviderRequest",
+    "RubricProvider",
+    "get_provider",
+    "register_provider",
+    "supported_providers",
+]

osmosis_ai/providers/anthropic_provider.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+try:  # pragma: no cover - optional dependency
+    import anthropic  # type: ignore
+    from anthropic import APIError  # type: ignore
+except ImportError:  # pragma: no cover - optional dependency
+    anthropic = None  # type: ignore[assignment]
+    APIError = None  # type: ignore[assignment]
+
+from ..rubric_types import ModelNotFoundError, ProviderRequestError, RewardRubricRunResult
+from .base import DEFAULT_REQUEST_TIMEOUT_SECONDS, ProviderRequest, RubricProvider
+from .shared import dump_model, extract_structured_score, reward_schema_definition
+
+
+class AnthropicProvider(RubricProvider):
+    name = "anthropic"
+
+    def default_timeout(self, model: str) -> float:
+        return DEFAULT_REQUEST_TIMEOUT_SECONDS
+
+    def run(self, request: ProviderRequest) -> RewardRubricRunResult:
+        if anthropic is None or APIError is None:
+            raise ProviderRequestError(
+                self.name,
+                request.model,
+                "Anthropic SDK is required. Install it via `pip install anthropic`.",
+            )
+
+        client = anthropic.Anthropic(api_key=request.api_key)
+        tool_name = "emit_reward_rubric_response"
+        schema_definition = reward_schema_definition()
+        tool = {
+            "name": tool_name,
+            "description": "Return the reward rubric score and explanation as structured JSON.",
+            "input_schema": schema_definition,
+        }
+
+        try:
+            response = client.messages.create(
+                model=request.model,
+                system=request.system_content,
+                messages=[{"role": "user", "content": [{"type": "text", "text": request.user_content}]}],
+                tools=[tool],
+                tool_choice={"type": "tool", "name": tool_name},
+                max_tokens=512,
+                temperature=0,
+                timeout=request.timeout,
+            )
+        except APIError as err:
+            detail = getattr(err, "message", None)
+            if not isinstance(detail, str) or not detail.strip():
+                detail = str(err)
+            status_code = getattr(err, "status_code", None)
+            if status_code == 404:
+                not_found_detail = (
+                    f"Model '{request.model}' was not found. Confirm your Anthropic account has access "
+                    "to the requested snapshot or update the model identifier."
+                )
+                raise ModelNotFoundError(self.name, request.model, not_found_detail) from err
+            raise ProviderRequestError(self.name, request.model, detail) from err
+        except Exception as err:
+            detail = str(err).strip() or "Unexpected error during Anthropic request."
+            raise ProviderRequestError(self.name, request.model, detail) from err
+
+        raw = dump_model(response)
+
+        payload: Dict[str, Any] | None = None
+        content_blocks = raw.get("content") if isinstance(raw, dict) else None
+        if isinstance(content_blocks, list):
+            for block in content_blocks:
+                if isinstance(block, dict) and block.get("type") == "tool_use" and block.get("name") == tool_name:
+                    maybe_input = block.get("input")
+                    if isinstance(maybe_input, dict):
+                        payload = maybe_input
+                    break
+        if payload is None:
+            raise ProviderRequestError(self.name, request.model, "Model response missing expected tool output.")
+        score, explanation = extract_structured_score(payload)
+        bounded = max(request.score_min, min(request.score_max, score))
+        return {"score": bounded, "explanation": explanation, "raw": raw}
+
+
+__all__ = ["AnthropicProvider"]

osmosis_ai/providers/base.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Tuple
+
+from ..rubric_types import RewardRubricRunResult
+
+DEFAULT_REQUEST_TIMEOUT_SECONDS = 30.0
+
+
+@dataclass(frozen=True)
+class ProviderRequest:
+    provider: str
+    model: str
+    api_key: str
+    system_content: str
+    user_content: str
+    score_min: float
+    score_max: float
+    timeout: float
+
+
+class RubricProvider:
+    """Interface for hosted LLM providers that can score rubrics."""
+
+    name: str
+
+    def default_timeout(self, model: str) -> float:
+        return DEFAULT_REQUEST_TIMEOUT_SECONDS
+
+    def run(self, request: ProviderRequest) -> RewardRubricRunResult:
+        raise NotImplementedError
+
+
+class ProviderRegistry:
+    def __init__(self) -> None:
+        self._providers: Dict[str, RubricProvider] = {}
+
+    def register(self, provider: RubricProvider) -> None:
+        key = provider.name
+        if key in self._providers:
+            raise ValueError(f"Provider '{key}' is already registered.")
+        self._providers[key] = provider
+
+    def get(self, name: str) -> RubricProvider:
+        try:
+            return self._providers[name]
+        except KeyError as exc:
+            raise ValueError(f"Unsupported provider '{name}'.") from exc
+
+    def supported_providers(self) -> Tuple[str, ...]:
+        return tuple(sorted(self._providers))
+
+
+__all__ = [
+    "DEFAULT_REQUEST_TIMEOUT_SECONDS",
+    "ProviderRequest",
+    "RubricProvider",
+    "ProviderRegistry",
+]