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

osmosis_ai/providers/gemini_provider.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+import inspect
 import time
 import warnings
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Tuple
@@ -167,6 +168,14 @@ def _seconds_to_millis(seconds: float) -> int:
     return max(int(round(seconds * 1000)), 1)
 
 
+def _supports_request_options(generate_content: Any) -> bool:
+    try:
+        signature = inspect.signature(generate_content)
+    except (TypeError, ValueError):
+        return False
+    return "request_options" in signature.parameters
+
+
 class GeminiProvider(RubricProvider):
     name = "gemini"
 
@@ -188,11 +197,31 @@ class GeminiProvider(RubricProvider):
         retry_timeouts = _build_retry_timeouts(requested_timeout)
         max_timeout = max(retry_timeouts)
 
+        supports_request_options = False
+        shared_client: Any | None = None
+
         with _suppress_pydantic_any_warning():
-            client = genai.Client(
+            probe_client = genai.Client(
                 api_key=request.api_key,
                 http_options={"timeout": _seconds_to_millis(max_timeout)},
             )
+        try:
+            supports_request_options = _supports_request_options(probe_client.models.generate_content)
+        except Exception:
+            try:
+                probe_client.close()
+            except Exception:
+                pass
+            raise
+
+        if supports_request_options:
+            shared_client = probe_client
+        else:
+            try:
+                probe_client.close()
+            except Exception:
+                pass
+
         schema_definition = reward_schema_definition()
         gemini_schema = _json_schema_to_genai(schema_definition, genai_types)
         config = genai_types.GenerateContentConfig(
@@ -206,33 +235,49 @@ class GeminiProvider(RubricProvider):
         response: Any | None = None
         last_error: Exception | None = None
 
-        for attempt_index, attempt_timeout in enumerate(retry_timeouts, start=1):
-            try:
-                with _suppress_pydantic_any_warning():
-                    try:
-                        response = client.models.generate_content(
-                            model=_normalize_gemini_model(request.model),
-                            contents=combined_prompt,
-                            config=config,
-                            request_options={"timeout": _seconds_to_millis(attempt_timeout)},
-                        )
-                    except TypeError as err:
-                        # Older SDKs may not accept request_options; retry without it.
-                        if "request_options" not in str(err):
-                            raise
-                        response = client.models.generate_content(
-                            model=_normalize_gemini_model(request.model),
-                            contents=combined_prompt,
-                            config=config,
-                        )
-                break
-            except Exception as err:  # pragma: no cover - network failures depend on runtime
-                last_error = err
-                if attempt_index >= len(retry_timeouts):
-                    detail = str(err).strip() or "Gemini request failed."
-                    raise ProviderRequestError(self.name, request.model, detail) from err
-                sleep_idx = min(attempt_index - 1, len(GEMINI_RETRY_SLEEP_SECONDS) - 1)
-                time.sleep(GEMINI_RETRY_SLEEP_SECONDS[sleep_idx])
+        try:
+            for attempt_index, attempt_timeout in enumerate(retry_timeouts, start=1):
+                per_attempt_client: Any | None = None
+                http_timeout_ms = _seconds_to_millis(attempt_timeout)
+                try:
+                    call_kwargs = {
+                        "model": _normalize_gemini_model(request.model),
+                        "contents": combined_prompt,
+                        "config": config,
+                    }
+                    if supports_request_options and shared_client is not None:
+                        call_client = shared_client
+                        call_kwargs["request_options"] = {"timeout": http_timeout_ms}
+                    else:
+                        with _suppress_pydantic_any_warning():
+                            per_attempt_client = genai.Client(
+                                api_key=request.api_key,
+                                http_options={"timeout": http_timeout_ms},
+                            )
+                        call_client = per_attempt_client
+
+                    with _suppress_pydantic_any_warning():
+                        response = call_client.models.generate_content(**call_kwargs)
+                    break
+                except Exception as err:  # pragma: no cover - network failures depend on runtime
+                    last_error = err
+                    if attempt_index >= len(retry_timeouts):
+                        detail = str(err).strip() or "Gemini request failed."
+                        raise ProviderRequestError(self.name, request.model, detail) from err
+                    sleep_idx = min(attempt_index - 1, len(GEMINI_RETRY_SLEEP_SECONDS) - 1)
+                    time.sleep(GEMINI_RETRY_SLEEP_SECONDS[sleep_idx])
+                finally:
+                    if per_attempt_client is not None:
+                        try:
+                            per_attempt_client.close()
+                        except Exception:
+                            pass
+        finally:
+            if shared_client is not None:
+                try:
+                    shared_client.close()
+                except Exception:
+                    pass
 
         if response is None and last_error is not None:
             detail = str(last_error).strip() or "Gemini request failed."

osmosis_ai/rubric_eval.py

@@ -206,74 +206,13 @@ def _build_user_prompt(
 
 
 def _collect_text_from_message(message: Dict[str, Any]) -> str:
+    from .cli_services.shared import collect_text_fragments
+
     content = message.get("content")
     if not isinstance(content, list):
         return ""
-    texts: List[str] = []
-
-    def _append_text(value: str) -> None:
-        stripped = value.strip()
-        if stripped:
-            texts.append(stripped)
-
-    def _walk(node: Any) -> None:
-        if isinstance(node, str):
-            _append_text(node)
-            return
-
-        if isinstance(node, list):
-            for item in node:
-                _walk(item)
-            return
-
-        if isinstance(node, dict):
-            # Prioritise common OpenAI / tool shapes, only escalating if a prior key yielded no text.
-            for key in ("text", "value"):
-                if key not in node:
-                    continue
-                before_count = len(texts)
-                _walk(node[key])
-                if len(texts) > before_count:
-                    break
-            if node.get("type") == "tool_result" and "content" in node:
-                _walk(node["content"])
-            elif "content" in node:
-                _walk(node["content"])
-            # Additional fallbacks (e.g., message wrappers).
-            for key in ("message", "parts", "input_text", "output_text"):
-                if key in node:
-                    _walk(node[key])
-            # Inspect remaining nested structures without re-traversing handled keys.
-            handled = {
-                "text",
-                "value",
-                "content",
-                "message",
-                "parts",
-                "input_text",
-                "output_text",
-                "type",
-                "role",
-                "name",
-                "id",
-                "index",
-                "finish_reason",
-                "reason",
-                "tool_call_id",
-                "metadata",
-            }
-            for key, value in node.items():
-                if key in handled:
-                    continue
-                if isinstance(value, (list, dict)):
-                    _walk(value)
-                elif isinstance(value, str) and key.lower() in {"text", "value", "message"}:
-                    _append_text(value)
-
-    for block in content:
-        _walk(block)
-
-    return " ".join(texts)
+    fragments = collect_text_fragments(content, allow_free_strings=True)
+    return " ".join(fragments)
 
 
 def _extract_latest_text(messages: List[Dict[str, Any]], role: str) -> Optional[str]:
@@ -386,6 +325,22 @@ def _resolve_api_key(provider: str, model_info: ModelInfo) -> str:
     return api_key
 
 
+def ensure_api_key_available(model_info: ModelInfo) -> None:
+    """
+    Validate that the provider specified in `model_info` has an accessible API key.
+
+    Raises:
+        MissingAPIKeyError: When the lookup fails or the environment variable is unset.
+        TypeError: When `model_info` is missing required fields.
+    """
+    provider_raw = model_info.get("provider")
+    if not isinstance(provider_raw, str) or not provider_raw.strip():
+        raise TypeError("'model_info' must include a 'provider' string")
+
+    provider = provider_raw.strip().lower()
+    _resolve_api_key(provider, model_info)
+
+
 def _run_reward_rubric(
     provider_name: str,
     provider_impl: RubricProvider,
@@ -534,4 +489,10 @@ def evaluate_rubric(
     return result if return_details else result["score"]
 
 
-__all__ = ["evaluate_rubric", "ModelInfo", "RewardRubricRunResult", "MissingAPIKeyError"]
+__all__ = [
+    "evaluate_rubric",
+    "ensure_api_key_available",
+    "ModelInfo",
+    "RewardRubricRunResult",
+    "MissingAPIKeyError",
+]

osmosis_ai/utils.py

@@ -28,10 +28,6 @@ def osmosis_reward(func: Callable) -> Callable:
     sig = inspect.signature(func)
     params = list(sig.parameters.values())
 
-    # Check parameter count
-    if len(params) < 2 or len(params) > 3:
-        raise TypeError(f"Function {func.__name__} must have 2-3 parameters, got {len(params)}")
-
     # Check first parameter: solution_str: str
     if params[0].name != 'solution_str':
         raise TypeError(f"First parameter must be named 'solution_str', got '{params[0].name}'")

{osmosis_ai-0.2.2.dist-info → osmosis_ai-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: osmosis-ai
-Version: 0.2.2
+Version: 0.2.3
 Summary: A Python library for reward function validation with strict type enforcement.
 Author-email: Osmosis AI <jake@osmosis.ai>
 License: MIT License
@@ -29,9 +29,18 @@ Project-URL: Issues, https://github.com/Osmosis-AI/osmosis-sdk-python/issues
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.6
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: PyYAML<7.0,>=6.0
+Requires-Dist: python-dotenv<2.0.0,>=0.1.0
+Requires-Dist: requests<3.0.0,>=2.0.0
+Requires-Dist: xxhash<4.0.0,>=3.0.0
+Requires-Dist: anthropic<0.50.0,>=0.36.0
+Requires-Dist: openai>=2.0.0
+Requires-Dist: google-genai>=1.0.0
+Requires-Dist: xai-sdk>=1.2.0
+Requires-Dist: tqdm<5.0.0,>=4.0.0
 Dynamic: license-file
 
 # osmosis-ai
@@ -44,6 +53,10 @@ A Python library that provides reward and rubric validation helpers for LLM appl
 pip install osmosis-ai
 ```
 
+Requires Python 3.9 or newer.
+
+This installs the Osmosis CLI and pulls in the required provider SDKs (`openai`, `anthropic`, `google-genai`, `xai-sdk`) along with supporting utilities such as `PyYAML`, `python-dotenv`, `requests`, and `xxhash`.
+
 For development:
 ```bash
 git clone https://github.com/Osmosis-AI/osmosis-sdk-python
@@ -211,6 +224,55 @@ def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = N
 
 - `examples/rubric_functions.py` demonstrates `evaluate_rubric` with OpenAI, Anthropic, Gemini, and xAI using the schema-enforced SDK integrations.
 - `examples/reward_functions.py` keeps local reward helpers that showcase the decorator contract without external calls.
+- `examples/rubric_configs.yaml` bundles two rubric definitions, each with its own provider configuration and extra prompt context.
+- `examples/sample_data.jsonl` contains two conversation payloads mapped to those rubrics so you can trial dataset validation.
+
+```yaml
+# examples/rubric_configs.yaml (excerpt)
+version: 1
+rubrics:
+  - id: support_followup
+    model_info:
+      provider: openai
+      model: gpt-5-mini
+      api_key_env: OPENAI_API_KEY
+```
+
+```jsonl
+{"conversation_id": "ticket-001", "rubric_id": "support_followup", "...": "..."}
+{"conversation_id": "ticket-047", "rubric_id": "policy_grounding", "...": "..."}
+```
+
+## CLI Tools
+
+Installing the SDK also provides a lightweight CLI available as `osmosis` (aliases: `osmosis_ai`, `osmosis-ai`) for inspecting rubric YAML files and JSONL test payloads.
+
+Preview a rubric file and print every configuration discovered, including nested entries:
+
+```bash
+osmosis preview --path path/to/rubric.yaml
+```
+
+Preview a dataset of chat transcripts stored as JSONL:
+
+```bash
+osmosis preview --path path/to/data.jsonl
+```
+
+Evaluate a dataset against a hosted rubric configuration and print the returned scores:
+
+```bash
+osmosis eval --rubric support_followup --data examples/sample_data.jsonl
+```
+
+- Supply the dataset with `-d`/`--data path/to/data.jsonl`; the path is resolved relative to the current working directory.
+- Use `--config path/to/rubric_configs.yaml` when the rubric definitions are not located alongside the dataset.
+- Pass `-n`/`--number` to sample the provider multiple times per record; the CLI prints every run along with aggregate statistics (average, variance, standard deviation, and min/max).
+- Provide `--output path/to/dir` to create the directory (if needed) and emit `rubric_eval_result_<unix_timestamp>.json`, or supply a full file path (any extension) to control the filename; each file captures every run, provider payloads, timestamps, and aggregate statistics for downstream analysis.
+- Skip `--output` to collect results under `~/.cache/osmosis/eval_result/<rubric_id>/rubric_eval_result_<identifier>.json`; the CLI writes this JSON whether the evaluation finishes cleanly or hits provider/runtime errors so you can inspect failures later (only a manual Ctrl+C interrupt leaves no file behind).
+- Dataset rows whose `rubric_id` does not match the requested rubric are skipped automatically.
+
+Both commands validate the file, echo a short summary (`Loaded <n> ...`), and pretty-print the parsed records so you can confirm that new rubrics or test fixtures look correct before committing them. Invalid files raise a descriptive error and exit with a non-zero status code.
 
 ## Running Examples
 

osmosis_ai-0.2.3.dist-info/RECORD

@@ -0,0 +1,27 @@
+osmosis_ai/__init__.py,sha256=2_qXxu18Yc7UicqxFZds8PjR4q0mTY1Xt17iR38OFbw,725
+osmosis_ai/cli.py,sha256=EPCttBnj1TEqQuO2gmS9iHadYcudiizVM38jACztRFE,1320
+osmosis_ai/cli_commands.py,sha256=CmTcb5N3delW7z3fwucss89xw5MHgIrJJ2Z5xdAuIeU,6165
+osmosis_ai/consts.py,sha256=-NDo9FaqBTebkCnhiFDxne6BY0W7BL3oM8HnGQDDgSE,73
+osmosis_ai/rubric_eval.py,sha256=PE2MvJygMbxelsJSTRzlW0bf-YUrtc8lCh6iTpHkjnU,17029
+osmosis_ai/rubric_types.py,sha256=kJvNAjLd3Y-1Q-_Re9HLTprLAUO3qtwR-IWOBeMkFI8,1279
+osmosis_ai/utils.py,sha256=IfTicRfa2Ybut4OzV4pHGSLBv-sGcmdT4eKIrIq4Pj8,19758
+osmosis_ai/cli_services/__init__.py,sha256=QQBwlI4KXoXK1X_e7kwW5sAVSh1VqBVuPllCwUOGXDM,1534
+osmosis_ai/cli_services/config.py,sha256=5hW2taAMhO9BkOfXvUCclnKLKtGTPaytS0oAgxCqymY,13965
+osmosis_ai/cli_services/dataset.py,sha256=qA0WHuOlJZCZdDbFX7ltimaZ3ujZpVUah4pxMvd4lVk,9042
+osmosis_ai/cli_services/engine.py,sha256=DbdJ24e5njk_lihe00cUgLO7yyJ6YgekFy6MAg_uq0k,9157
+osmosis_ai/cli_services/errors.py,sha256=nI6jlICyA4MMNKmwDHQBwyJVah5PVwstmra1HpGkVLE,136
+osmosis_ai/cli_services/reporting.py,sha256=H2g0BmEE2stVey4RmurQM713VowH8984a9r7oDstSkA,12499
+osmosis_ai/cli_services/session.py,sha256=Ru3HA80eqRYZGD1e38N8yd96FiAY8cIYpJvEOHKakM0,6597
+osmosis_ai/cli_services/shared.py,sha256=PilPfW5oDvNL5VG8oObSq2ZL35QPFmhBDf0V4gfd2Ro,5942
+osmosis_ai/providers/__init__.py,sha256=yLSExLbJToZ8AUOVxt4LDplxtIuwv-etSJJyZOcOE2Q,927
+osmosis_ai/providers/anthropic_provider.py,sha256=zrWCVP8co4v8xhcJDFLASwvwEADKN-1p34cY_GH4q5M,3758
+osmosis_ai/providers/base.py,sha256=fN5cnWXYAHN53RR_x6ykbUkM4bictNPDj4U8yd4b2a0,1492
+osmosis_ai/providers/gemini_provider.py,sha256=QANSCmkKungpkpDP2RClmKYnwNVrGv3MKxJwkh68IhY,12045
+osmosis_ai/providers/openai_family.py,sha256=DeQWPMcafEvG4xcI97m3AADTKP2pYw9KwcQTcQg-h_4,26078
+osmosis_ai/providers/shared.py,sha256=dmVe8JDgafPmo6HkP-Kl0aWfffhAT6u3ElV_wLlYD34,2957
+osmosis_ai-0.2.3.dist-info/licenses/LICENSE,sha256=FV2ZmyhdCYinoLLvU_ci-7pZ3DeNYY9XqZjVjOd3h94,1064
+osmosis_ai-0.2.3.dist-info/METADATA,sha256=tgRpinJ60KxX1OlT_9JMCFl2hY0A5I4rZLPRtXQ4p-U,13669
+osmosis_ai-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+osmosis_ai-0.2.3.dist-info/entry_points.txt,sha256=aF1CR36a9I9_vcF7nlK9JnK1Iqu614vPy2_jh4QU26A,114
+osmosis_ai-0.2.3.dist-info/top_level.txt,sha256=UPNRTKIBSrxsJVNxwXnLCqSoBS4bAiL_3jMtjvf5zEY,11
+osmosis_ai-0.2.3.dist-info/RECORD,,

osmosis_ai-0.2.3.dist-info/entry_points.txt

@@ -0,0 +1,4 @@
+[console_scripts]
+osmosis = osmosis_ai.cli:main
+osmosis-ai = osmosis_ai.cli:main
+osmosis_ai = osmosis_ai.cli:main

osmosis_ai-0.2.2.dist-info/RECORD

@@ -1,16 +0,0 @@
-osmosis_ai/__init__.py,sha256=2_qXxu18Yc7UicqxFZds8PjR4q0mTY1Xt17iR38OFbw,725
-osmosis_ai/consts.py,sha256=-NDo9FaqBTebkCnhiFDxne6BY0W7BL3oM8HnGQDDgSE,73
-osmosis_ai/rubric_eval.py,sha256=bFgxgnbQeD-7K2LkTJfnSk5aG9s4lefLfmvQt4GQSnM,18332
-osmosis_ai/rubric_types.py,sha256=kJvNAjLd3Y-1Q-_Re9HLTprLAUO3qtwR-IWOBeMkFI8,1279
-osmosis_ai/utils.py,sha256=yjC_oQt1wwTJsX7lCx0ZGMa5txHURByuBDuU37WPAO0,19927
-osmosis_ai/providers/__init__.py,sha256=yLSExLbJToZ8AUOVxt4LDplxtIuwv-etSJJyZOcOE2Q,927
-osmosis_ai/providers/anthropic_provider.py,sha256=zrWCVP8co4v8xhcJDFLASwvwEADKN-1p34cY_GH4q5M,3758
-osmosis_ai/providers/base.py,sha256=fN5cnWXYAHN53RR_x6ykbUkM4bictNPDj4U8yd4b2a0,1492
-osmosis_ai/providers/gemini_provider.py,sha256=xqklXRO5K1YZ4SKq5lfU3bDUaF8QN2MIBP4DHGKwLVo,10611
-osmosis_ai/providers/openai_family.py,sha256=DeQWPMcafEvG4xcI97m3AADTKP2pYw9KwcQTcQg-h_4,26078
-osmosis_ai/providers/shared.py,sha256=dmVe8JDgafPmo6HkP-Kl0aWfffhAT6u3ElV_wLlYD34,2957
-osmosis_ai-0.2.2.dist-info/licenses/LICENSE,sha256=FV2ZmyhdCYinoLLvU_ci-7pZ3DeNYY9XqZjVjOd3h94,1064
-osmosis_ai-0.2.2.dist-info/METADATA,sha256=MPovk4NSQ_viTMd-zx1lp7Uo2EGB3EotjftcNddy4HU,10448
-osmosis_ai-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-osmosis_ai-0.2.2.dist-info/top_level.txt,sha256=UPNRTKIBSrxsJVNxwXnLCqSoBS4bAiL_3jMtjvf5zEY,11
-osmosis_ai-0.2.2.dist-info/RECORD,,