osmosis-ai 0.1.9__py3-none-any.whl → 0.2.4__py3-none-any.whl

osmosis_ai/__init__.py

@@ -1,137 +1,24 @@
 """
-osmosis-ai: A library for monkey patching LLM APIs to print all prompts and responses.
+osmosis-ai: A Python library for reward function validation with strict type enforcement.
 
-This module patches various LLM client libraries to send all prompts and responses
-to the OSMOSIS API for logging and monitoring.
+This library provides decorators such as @osmosis_reward and @osmosis_rubric that
+enforce standardized function signatures for LLM-centric workflows.
 
-Currently supported adapters:
-- Anthropic (both sync and async clients)
-- OpenAI (both sync and async clients, v1 and v2 API versions)
-- LangChain (LLMs, Chat Models, and Prompt Templates)
+Features:
+- Type-safe reward function decoration
+- Parameter name and type validation
+- Support for optional configuration parameters
 """
 
-
-# Use lazy imports to avoid importing modules during setup
-def _import_modules():
-    global utils, logger, reconfigure_logger
-    global set_log_destination, log_destination, LogDestination
-    global init, enabled, disable_osmosis, enable_osmosis, osmosis_reward
-
-    from . import utils
-    from .logger import logger, reconfigure_logger, set_log_destination, log_destination
-    from .consts import LogDestination
-
-    # Re-export configuration flags for easy access
-    enabled = utils.enabled
-
-    # Export disable and enable functions
-    disable_osmosis = utils.disable_osmosis
-    enable_osmosis = utils.enable_osmosis
-
-    # Re-export initialization function
-    init = utils.init
-    
-    # Export the reward decorator
-    osmosis_reward = utils.osmosis_reward
-
-    # Initialize wrappers as None
-    global wrap_anthropic, wrap_openai, wrap_langchain
-    global wrap_langchain_openai, wrap_langchain_anthropic
-
-    wrap_anthropic = None
-    wrap_openai = None
-    wrap_langchain = None
-    wrap_langchain_openai = None
-    wrap_langchain_anthropic = None
-
-    # Lazily set up anthropic wrapper if available
-    try:
-        from .adapters import anthropic
-
-        wrap_anthropic = anthropic.wrap_anthropic
-        # Apply the wrapper, but don't fail if it doesn't work
-        try:
-            wrap_anthropic()
-        except Exception as e:
-            logger.warning(f"Failed to wrap Anthropic: {str(e)}")
-    except ImportError:
-
-        def wrap_anthropic():
-            logger.warning(
-                "Anthropic support not available. Install Anthropic to use this feature."
-            )
-
-    # Lazily set up OpenAI wrapper if available
-    try:
-        from .adapters import openai
-
-        wrap_openai = openai.wrap_openai
-        # Apply the wrapper, but don't fail if it doesn't work
-        try:
-            wrap_openai()
-        except Exception as e:
-            logger.warning(f"Failed to wrap OpenAI: {str(e)}")
-    except ImportError:
-
-        def wrap_openai():
-            logger.warning(
-                "OpenAI support not available. Install OpenAI to use this feature."
-            )
-
-    # Lazily set up LangChain wrapper if available
-    try:
-        from .adapters import langchain
-
-        wrap_langchain = langchain.wrap_langchain
-        # Apply the wrapper, but don't fail if it doesn't work
-        try:
-            wrap_langchain()
-        except Exception as e:
-            logger.warning(f"Failed to wrap LangChain: {str(e)}")
-    except ImportError:
-
-        def wrap_langchain():
-            logger.warning(
-                "LangChain support not available. Install LangChain to use this feature."
-            )
-
-    # Lazily set up LangChain-OpenAI wrapper if available
-    try:
-        from .adapters import langchain_openai
-
-        wrap_langchain_openai = langchain_openai.wrap_langchain_openai
-        # Apply the wrapper, but don't fail if it doesn't work
-        try:
-            wrap_langchain_openai()
-        except Exception as e:
-            logger.warning(f"Failed to wrap LangChain-OpenAI: {str(e)}")
-    except ImportError:
-
-        def wrap_langchain_openai():
-            logger.warning(
-                "LangChain-OpenAI support not available. Install LangChain and OpenAI to use this feature."
-            )
-
-    # Lazily set up LangChain-Anthropic wrapper if available
-    try:
-        from .adapters import langchain_anthropic
-
-        wrap_langchain_anthropic = langchain_anthropic.wrap_langchain_anthropic
-        # Apply the wrapper, but don't fail if it doesn't work
-        try:
-            wrap_langchain_anthropic()
-        except Exception as e:
-            logger.warning(f"Failed to wrap LangChain-Anthropic: {str(e)}")
-    except ImportError:
-
-        def wrap_langchain_anthropic():
-            logger.warning(
-                "LangChain-Anthropic support not available. Install LangChain and Anthropic to use this feature."
-            )
-
-
-# Initialize the module on first import, but not during installation
-import sys
-
-if "pip" not in sys.modules:
-    _import_modules()
+from .rubric_eval import MissingAPIKeyError, evaluate_rubric
+from .rubric_types import ModelNotFoundError, ProviderRequestError
+from .utils import osmosis_reward, osmosis_rubric
+
+__all__ = [
+    "osmosis_reward",
+    "osmosis_rubric",
+    "evaluate_rubric",
+    "MissingAPIKeyError",
+    "ProviderRequestError",
+    "ModelNotFoundError",
+]

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,60 @@
+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 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",
+    "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",
+]