From 37279872ebb62fadc04f8bdaa5311918cdf95ce6 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 10:56:24 -0800
Subject: [PATCH 01/60] add ruff rules
---
examples/reward_functions.py | 38 +-
examples/rubric_functions.py | 22 +-
osmosis_ai/__init__.py | 80 +-
osmosis_ai/auth/__init__.py | 16 +-
osmosis_ai/auth/credentials.py | 24 +-
osmosis_ai/auth/flow.py | 31 +-
osmosis_ai/auth/local_server.py | 17 +-
osmosis_ai/auth/platform_client.py | 24 +-
osmosis_ai/cli.py | 19 +-
osmosis_ai/cli_commands.py | 92 +-
osmosis_ai/cli_services/__init__.py | 19 +-
osmosis_ai/cli_services/config.py | 69 +-
osmosis_ai/cli_services/dataset.py | 65 +-
osmosis_ai/cli_services/engine.py | 111 +-
osmosis_ai/cli_services/reporting.py | 105 +-
osmosis_ai/cli_services/session.py | 68 +-
osmosis_ai/cli_services/shared.py | 33 +-
osmosis_ai/rollout/__init__.py | 204 ++--
osmosis_ai/rollout/_compat.py | 20 +-
osmosis_ai/rollout/cli.py | 28 +-
osmosis_ai/rollout/cli_utils.py | 8 +-
osmosis_ai/rollout/client.py | 60 +-
osmosis_ai/rollout/config/__init__.py | 6 +-
osmosis_ai/rollout/config/settings.py | 8 +-
osmosis_ai/rollout/console.py | 28 +-
osmosis_ai/rollout/core/__init__.py | 66 +-
osmosis_ai/rollout/core/base.py | 34 +-
osmosis_ai/rollout/core/exceptions.py | 8 +-
osmosis_ai/rollout/core/llm_client.py | 8 +-
osmosis_ai/rollout/core/schemas.py | 103 +-
osmosis_ai/rollout/eval/common/__init__.py | 24 +-
osmosis_ai/rollout/eval/common/cli.py | 57 +-
osmosis_ai/rollout/eval/common/dataset.py | 55 +-
osmosis_ai/rollout/eval/common/errors.py | 7 +-
osmosis_ai/rollout/eval/common/llm_client.py | 31 +-
osmosis_ai/rollout/eval/common/runner.py | 54 +-
.../rollout/eval/evaluation/__init__.py | 14 +-
osmosis_ai/rollout/eval/evaluation/cli.py | 42 +-
osmosis_ai/rollout/eval/evaluation/eval_fn.py | 24 +-
osmosis_ai/rollout/eval/evaluation/report.py | 24 +-
osmosis_ai/rollout/eval/evaluation/runner.py | 189 ++--
osmosis_ai/rollout/eval/test_mode/__init__.py | 14 +-
osmosis_ai/rollout/eval/test_mode/cli.py | 58 +-
.../rollout/eval/test_mode/interactive.py | 101 +-
osmosis_ai/rollout/eval/test_mode/runner.py | 30 +-
osmosis_ai/rollout/mcp/agent_loop.py | 30 +-
osmosis_ai/rollout/mcp/loader.py | 6 +-
osmosis_ai/rollout/network.py | 29 +-
osmosis_ai/rollout/registry.py | 16 +-
osmosis_ai/rollout/server/__init__.py | 12 +-
osmosis_ai/rollout/server/api_key.py | 7 +-
osmosis_ai/rollout/server/app.py | 60 +-
osmosis_ai/rollout/server/registration.py | 30 +-
osmosis_ai/rollout/server/serve.py | 28 +-
osmosis_ai/rollout/server/state.py | 38 +-
osmosis_ai/rollout/testing.py | 48 +-
osmosis_ai/rollout/tools.py | 31 +-
osmosis_ai/rollout/utils.py | 26 +-
osmosis_ai/rollout/validator.py | 47 +-
osmosis_ai/rubric_eval.py | 124 +-
osmosis_ai/rubric_types.py | 24 +-
osmosis_ai/utils.py | 109 +-
pyproject.toml | 28 +-
tests/conftest.py | 16 +-
tests/test_auth_credentials.py | 4 +-
tests/test_cli.py | 40 +-
tests/test_cli_services.py | 48 +-
tests/test_eval_fn.py | 4 +-
tests/test_eval_report.py | 2 +-
tests/test_eval_runner.py | 65 +-
tests/test_litellm_provider.py | 182 ++-
tests/test_mcp_agent_loop.py | 121 +-
tests/test_rollout_api_key.py | 8 +-
tests/test_rollout_base.py | 8 +-
tests/test_rollout_client.py | 6 +-
tests/test_rollout_registry.py | 10 +-
tests/test_rollout_schemas.py | 5 +-
tests/test_rollout_server.py | 39 +-
tests/test_rollout_testing.py | 7 +-
tests/test_rollout_tools.py | 3 +-
tests/test_rollout_utils.py | 9 +-
tests/test_rollout_validator.py | 17 +-
tests/test_test_mode_dataset.py | 22 +-
tests/test_test_mode_providers.py | 83 +-
tests/test_test_mode_runner.py | 31 +-
uv.lock | 1001 +++++++++++++++--
86 files changed, 3048 insertions(+), 1514 deletions(-)
diff --git a/examples/reward_functions.py b/examples/reward_functions.py
index af10af86..37160381 100644
--- a/examples/reward_functions.py
+++ b/examples/reward_functions.py
@@ -7,41 +7,51 @@
from osmosis_ai import osmosis_reward
-
# CORRECT USAGE EXAMPLES
+
@osmosis_reward
-def simple_exact_match(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+def simple_exact_match(
+ solution_str: str, ground_truth: str, extra_info: dict | None = None
+) -> float:
"""Basic exact match reward function."""
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
@osmosis_reward
-def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+def case_insensitive_match(
+ solution_str: str, ground_truth: str, extra_info: dict | None = None
+) -> float:
"""Case-insensitive string matching with optional extra info."""
match = solution_str.lower().strip() == ground_truth.lower().strip()
# Use extra_info if provided
- if extra_info and 'partial_credit' in extra_info:
- if not match and extra_info['partial_credit']:
- # Give partial credit for similar length
- len_diff = abs(len(solution_str) - len(ground_truth))
- if len_diff <= 2:
- return 0.5
+ if (
+ extra_info
+ and "partial_credit" in extra_info
+ and not match
+ and extra_info["partial_credit"]
+ ):
+ # Give partial credit for similar length
+ len_diff = abs(len(solution_str) - len(ground_truth))
+ if len_diff <= 2:
+ return 0.5
return 1.0 if match else 0.0
@osmosis_reward
-def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+def numeric_tolerance(
+ solution_str: str, ground_truth: str, extra_info: dict | None = None
+) -> float:
"""Numeric comparison with tolerance."""
try:
solution_num = float(solution_str.strip())
truth_num = float(ground_truth.strip())
tolerance = 0.01 # default
- if extra_info and 'tolerance' in extra_info:
- tolerance = extra_info['tolerance']
+ if extra_info and "tolerance" in extra_info:
+ tolerance = extra_info["tolerance"]
return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
except ValueError:
@@ -93,7 +103,7 @@ def minimal_reward(solution_str: str, ground_truth: str) -> float:
simple_exact_match,
case_insensitive_match,
numeric_tolerance,
- minimal_reward
+ minimal_reward,
]
for func in functions:
@@ -103,4 +113,4 @@ def minimal_reward(solution_str: str, ground_truth: str) -> float:
result = func(solution, truth, extra)
print(f" {solution!r} vs {truth!r} (extra={extra}) -> {result}")
except Exception as e:
- print(f" {solution!r} vs {truth!r} (extra={extra}) -> ERROR: {e}")
\ No newline at end of file
+ print(f" {solution!r} vs {truth!r} (extra={extra}) -> ERROR: {e}")
diff --git a/examples/rubric_functions.py b/examples/rubric_functions.py
index 03f3da6b..3d7a2346 100644
--- a/examples/rubric_functions.py
+++ b/examples/rubric_functions.py
@@ -102,11 +102,21 @@ def score_with_hosted_model(
Provide provider-specific knobs inside `extra_info["metadata"]`. Toggle `extra_info["capture_details"]`
when you want the provider response returned.
"""
- metadata = extra_info.get("metadata") if isinstance(extra_info, dict) and isinstance(extra_info.get("metadata"), dict) else None
- capture_details = bool(extra_info.get("capture_details")) if isinstance(extra_info, dict) else False
+ metadata = (
+ extra_info.get("metadata")
+ if isinstance(extra_info, dict) and isinstance(extra_info.get("metadata"), dict)
+ else None
+ )
+ capture_details = (
+ bool(extra_info.get("capture_details"))
+ if isinstance(extra_info, dict)
+ else False
+ )
prompt_metadata = metadata
- provider_config = _resolve_provider_profile(metadata.get("provider_profile") if metadata else None)
+ provider_config = _resolve_provider_profile(
+ metadata.get("provider_profile") if metadata else None
+ )
model_info = dict(provider_config["model_info"])
@@ -148,7 +158,9 @@ def _resolve_provider_profile(profile_name: str | None) -> dict:
profile = PROFILE_CATALOG.get(profile_key)
if profile is None:
options = ", ".join(sorted(PROFILE_CATALOG))
- raise ValueError(f"Unknown provider_profile '{profile_name}'. Supported profiles: {options}")
+ raise ValueError(
+ f"Unknown provider_profile '{profile_name}'. Supported profiles: {options}"
+ )
return {
"model_info": profile,
@@ -165,7 +177,7 @@ def _run(provider_name: str, provider_profile: str) -> None:
"metadata": {
"provider_profile": provider_profile,
"scenario_label": provider_name,
- }
+ },
}
score = score_with_hosted_model(
solution_str=SOLUTION_STR,
diff --git a/osmosis_ai/__init__.py b/osmosis_ai/__init__.py
index 7ba0a3b0..1e51ddd5 100644
--- a/osmosis_ai/__init__.py
+++ b/osmosis_ai/__init__.py
@@ -8,67 +8,67 @@
"""
from .consts import PACKAGE_VERSION as __version__
-from .rubric_eval import MissingAPIKeyError, evaluate_rubric
-from .rubric_types import ModelNotFoundError, ProviderRequestError
-from .utils import osmosis_reward, osmosis_rubric
# Remote rollout SDK exports
from .rollout import (
+ CompletionsResult,
+ InitResponse,
+ OpenAIFunctionToolSchema,
+ # Client
+ OsmosisLLMClient,
+ # Exceptions
+ OsmosisRolloutError,
+ OsmosisServerError,
+ OsmosisTimeoutError,
+ OsmosisTransportError,
+ OsmosisValidationError,
# Core classes
RolloutAgentLoop,
RolloutContext,
+ RolloutMetrics,
+ # Schemas
+ RolloutRequest,
+ RolloutResponse,
RolloutResult,
- # Client
- OsmosisLLMClient,
- CompletionsResult,
# Server
create_app,
- # Registry
- register_agent_loop,
get_agent_loop,
list_agent_loops,
- # Schemas
- RolloutRequest,
- RolloutResponse,
- InitResponse,
- OpenAIFunctionToolSchema,
- RolloutMetrics,
- # Exceptions
- OsmosisRolloutError,
- OsmosisTransportError,
- OsmosisServerError,
- OsmosisValidationError,
- OsmosisTimeoutError,
+ # Registry
+ register_agent_loop,
)
+from .rubric_eval import MissingAPIKeyError, evaluate_rubric
+from .rubric_types import ModelNotFoundError, ProviderRequestError
+from .utils import osmosis_reward, osmosis_rubric
__all__ = [
- # Version
- "__version__",
- # Reward function decorators
- "osmosis_reward",
- "osmosis_rubric",
- "evaluate_rubric",
+ "CompletionsResult",
+ "InitResponse",
"MissingAPIKeyError",
- "ProviderRequestError",
"ModelNotFoundError",
+ "OpenAIFunctionToolSchema",
+ "OsmosisLLMClient",
+ "OsmosisRolloutError",
+ "OsmosisServerError",
+ "OsmosisTimeoutError",
+ "OsmosisTransportError",
+ "OsmosisValidationError",
+ "ProviderRequestError",
# Remote rollout SDK
"RolloutAgentLoop",
"RolloutContext",
+ "RolloutMetrics",
+ "RolloutRequest",
+ "RolloutResponse",
"RolloutResult",
- "OsmosisLLMClient",
- "CompletionsResult",
+ # Version
+ "__version__",
"create_app",
- "register_agent_loop",
+ "evaluate_rubric",
"get_agent_loop",
"list_agent_loops",
- "RolloutRequest",
- "RolloutResponse",
- "InitResponse",
- "OpenAIFunctionToolSchema",
- "RolloutMetrics",
- "OsmosisRolloutError",
- "OsmosisTransportError",
- "OsmosisServerError",
- "OsmosisValidationError",
- "OsmosisTimeoutError",
+ # Reward function decorators
+ "osmosis_reward",
+ "osmosis_rubric",
+ "register_agent_loop",
]
diff --git a/osmosis_ai/auth/__init__.py b/osmosis_ai/auth/__init__.py
index da44acc3..59d0eb4e 100644
--- a/osmosis_ai/auth/__init__.py
+++ b/osmosis_ai/auth/__init__.py
@@ -27,9 +27,15 @@
"CONFIG_DIR",
"CREDENTIALS_FILE",
"PLATFORM_URL",
+ # Platform Client
+ "AuthenticationExpiredError",
# Credentials
"CredentialsStore",
+ # Flow
+ "LoginError",
+ "LoginResult",
"OrganizationInfo",
+ "PlatformAPIError",
"UserInfo",
"WorkspaceCredentials",
"delete_credentials",
@@ -38,14 +44,8 @@
"get_all_workspaces",
"get_valid_credentials",
"load_credentials",
- "save_credentials",
- "set_active_workspace",
- # Flow
- "LoginError",
- "LoginResult",
"login",
- # Platform Client
- "AuthenticationExpiredError",
- "PlatformAPIError",
"platform_request",
+ "save_credentials",
+ "set_active_workspace",
]
diff --git a/osmosis_ai/auth/credentials.py b/osmosis_ai/auth/credentials.py
index ae07ab90..fe127c4b 100644
--- a/osmosis_ai/auth/credentials.py
+++ b/osmosis_ai/auth/credentials.py
@@ -8,7 +8,7 @@
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
-from typing import Any, Optional
+from typing import Any
from .config import CONFIG_DIR, CREDENTIALS_FILE, CREDENTIALS_VERSION
@@ -19,7 +19,7 @@ class UserInfo:
id: str
email: str
- name: Optional[str] = None
+ name: str | None = None
def to_dict(self) -> dict[str, Any]:
return {"id": self.id, "email": self.email, "name": self.name}
@@ -63,7 +63,7 @@ class WorkspaceCredentials:
user: UserInfo
organization: OrganizationInfo
created_at: datetime
- token_id: Optional[str] = None # Platform token ID for revocation
+ token_id: str | None = None # Platform token ID for revocation
def to_dict(self) -> dict[str, Any]:
result = {
@@ -82,7 +82,9 @@ def to_dict(self) -> dict[str, Any]:
def from_dict(cls, data: dict[str, Any]) -> WorkspaceCredentials:
expires_at = datetime.fromisoformat(data["expires_at"])
if expires_at.tzinfo is None:
- raise ValueError("expires_at must be timezone-aware (ISO8601 with timezone offset)")
+ raise ValueError(
+ "expires_at must be timezone-aware (ISO8601 with timezone offset)"
+ )
return cls(
access_token=data["access_token"],
token_type=data["token_type"],
@@ -106,7 +108,7 @@ def is_expired(self) -> bool:
class CredentialsStore:
"""Multi-workspace credentials storage."""
- active_workspace: Optional[str] # workspace name
+ active_workspace: str | None # workspace name
workspaces: dict[str, WorkspaceCredentials] # keyed by workspace name
def to_dict(self) -> dict[str, Any]:
@@ -128,7 +130,7 @@ def from_dict(cls, data: dict[str, Any]) -> CredentialsStore:
workspaces=workspaces,
)
- def get_active_credentials(self) -> Optional[WorkspaceCredentials]:
+ def get_active_credentials(self) -> WorkspaceCredentials | None:
"""Get credentials for the active workspace."""
if not self.active_workspace:
return None
@@ -155,13 +157,13 @@ def _set_file_permissions(path: Path) -> None:
os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 0o600
-def _load_store() -> Optional[CredentialsStore]:
+def _load_store() -> CredentialsStore | None:
"""Load the credentials store from file."""
if not CREDENTIALS_FILE.exists():
return None
try:
- with open(CREDENTIALS_FILE, "r", encoding="utf-8") as f:
+ with open(CREDENTIALS_FILE, encoding="utf-8") as f:
data = json.load(f)
return CredentialsStore.from_dict(data)
except (json.JSONDecodeError, KeyError, ValueError):
@@ -205,7 +207,7 @@ def save_credentials(credentials: WorkspaceCredentials) -> None:
_save_store(store)
-def load_credentials() -> Optional[WorkspaceCredentials]:
+def load_credentials() -> WorkspaceCredentials | None:
"""Load credentials for the active workspace.
Returns:
@@ -255,7 +257,7 @@ def delete_workspace_credentials(workspace_name: str) -> bool:
return True
-def get_valid_credentials() -> Optional[WorkspaceCredentials]:
+def get_valid_credentials() -> WorkspaceCredentials | None:
"""Get credentials for the active workspace if they exist and are not expired.
Returns:
@@ -286,7 +288,7 @@ def get_all_workspaces() -> list[tuple[str, WorkspaceCredentials, bool]]:
return result
-def get_active_workspace() -> Optional[str]:
+def get_active_workspace() -> str | None:
"""Get the name of the active workspace.
Returns:
diff --git a/osmosis_ai/auth/flow.py b/osmosis_ai/auth/flow.py
index 74e11891..e65c1353 100644
--- a/osmosis_ai/auth/flow.py
+++ b/osmosis_ai/auth/flow.py
@@ -11,7 +11,12 @@
from ..consts import PACKAGE_VERSION
from .config import PLATFORM_URL
-from .credentials import WorkspaceCredentials, OrganizationInfo, UserInfo, save_credentials
+from .credentials import (
+ OrganizationInfo,
+ UserInfo,
+ WorkspaceCredentials,
+ save_credentials,
+)
from .local_server import LocalAuthServer, find_available_port
@@ -88,7 +93,7 @@ def login(
port = find_available_port()
if port is None:
raise LoginError(
- f"No available port found in range 8976-8985. "
+ "No available port found in range 8976-8985. "
"Please close any applications using these ports."
)
@@ -156,12 +161,16 @@ def login(
finally:
# Ensure the callback handler is unblocked if verification wasn't reached
if not server._verification_event.is_set():
- server.set_verification_result(success=False, error="Login failed unexpectedly")
+ server.set_verification_result(
+ success=False, error="Login failed unexpectedly"
+ )
server._shutdown_event.set()
server.server_close()
-def _verify_and_get_user_info(token: str) -> tuple[UserInfo, OrganizationInfo, datetime, str | None]:
+def _verify_and_get_user_info(
+ token: str,
+) -> tuple[UserInfo, OrganizationInfo, datetime, str | None]:
"""Verify token and get user info from the platform.
Args:
@@ -217,7 +226,9 @@ def _verify_and_get_user_info(token: str) -> tuple[UserInfo, OrganizationInfo, d
# Parse expiration - default to 90 days if not provided
expires_at_str = data.get("expires_at")
if expires_at_str:
- expires_at = datetime.fromisoformat(expires_at_str.replace("Z", "+00:00"))
+ expires_at = datetime.fromisoformat(
+ expires_at_str.replace("Z", "+00:00")
+ )
if expires_at.tzinfo is None:
raise LoginError(
"Invalid expires_at from platform: expected timezone-aware ISO8601 timestamp"
@@ -229,9 +240,9 @@ def _verify_and_get_user_info(token: str) -> tuple[UserInfo, OrganizationInfo, d
except HTTPError as e:
if e.code == 401:
- raise LoginError("Invalid or expired token")
- raise LoginError(f"Verification failed: HTTP {e.code}")
+ raise LoginError("Invalid or expired token") from None
+ raise LoginError(f"Verification failed: HTTP {e.code}") from e
except URLError as e:
- raise LoginError(f"Could not connect to platform: {e.reason}")
- except json.JSONDecodeError:
- raise LoginError("Invalid response from platform")
+ raise LoginError(f"Could not connect to platform: {e.reason}") from e
+ except json.JSONDecodeError as e:
+ raise LoginError("Invalid response from platform") from e
diff --git a/osmosis_ai/auth/local_server.py b/osmosis_ai/auth/local_server.py
index 2ae34e1c..b1497a31 100644
--- a/osmosis_ai/auth/local_server.py
+++ b/osmosis_ai/auth/local_server.py
@@ -6,7 +6,6 @@
import socket
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import Optional
from urllib.parse import parse_qs, urlparse
from .config import LOCAL_SERVER_PORT_END, LOCAL_SERVER_PORT_START
@@ -377,8 +376,8 @@ def __init__(self, port: int, expected_state: str) -> None:
"""
super().__init__(("localhost", port), AuthCallbackHandler)
self.expected_state = expected_state
- self.received_token: Optional[str] = None
- self.error: Optional[str] = None
+ self.received_token: str | None = None
+ self.error: str | None = None
self.revoked_count: int = 0
self._shutdown_event = threading.Event()
# Synchronization for deferred browser response
@@ -386,17 +385,21 @@ def __init__(self, port: int, expected_state: str) -> None:
self._verification_event = threading.Event()
self._verification_result: object = None # True = success, str = error message
- def set_verification_result(self, success: bool, error: Optional[str] = None) -> None:
+ def set_verification_result(self, success: bool, error: str | None = None) -> None:
"""Set the verification result and unblock the callback handler.
Args:
success: Whether token verification succeeded.
error: Error message if verification failed.
"""
- self._verification_result = True if success else (error or "Verification failed")
+ self._verification_result = (
+ True if success else (error or "Verification failed")
+ )
self._verification_event.set()
- def wait_for_callback(self, timeout: float = 300.0) -> tuple[Optional[str], Optional[str]]:
+ def wait_for_callback(
+ self, timeout: float = 300.0
+ ) -> tuple[str | None, str | None]:
"""Wait for the OAuth callback.
Args:
@@ -425,7 +428,7 @@ def shutdown(self) -> None:
super().shutdown()
-def find_available_port() -> Optional[int]:
+def find_available_port() -> int | None:
"""Find an available port in the configured range.
Returns:
diff --git a/osmosis_ai/auth/platform_client.py b/osmosis_ai/auth/platform_client.py
index 8ee1a2b3..5a5c6846 100644
--- a/osmosis_ai/auth/platform_client.py
+++ b/osmosis_ai/auth/platform_client.py
@@ -3,13 +3,17 @@
from __future__ import annotations
import json
-from typing import Any, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from ..consts import PACKAGE_VERSION
from .config import PLATFORM_URL
-from .credentials import delete_workspace_credentials, get_active_workspace, load_credentials
+from .credentials import (
+ delete_workspace_credentials,
+ get_active_workspace,
+ load_credentials,
+)
if TYPE_CHECKING:
from .credentials import WorkspaceCredentials
@@ -24,7 +28,7 @@ class AuthenticationExpiredError(Exception):
class PlatformAPIError(Exception):
"""Raised when a Platform API call fails."""
- def __init__(self, message: str, status_code: Optional[int] = None):
+ def __init__(self, message: str, status_code: int | None = None):
super().__init__(message)
self.status_code = status_code
@@ -44,10 +48,10 @@ def _handle_401_and_cleanup() -> None:
def platform_request(
endpoint: str,
method: str = "GET",
- data: Optional[dict[str, Any]] = None,
- headers: Optional[dict[str, str]] = None,
+ data: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
timeout: float = 30.0,
- credentials: Optional["WorkspaceCredentials"] = None,
+ credentials: WorkspaceCredentials | None = None,
) -> dict[str, Any]:
"""Make an authenticated request to the Osmosis Platform API.
@@ -108,8 +112,8 @@ def platform_request(
# Ignore body read/decoding failures; keep the original status code context.
pass
- raise PlatformAPIError(f"API error: HTTP {e.code}.{detail}", e.code)
+ raise PlatformAPIError(f"API error: HTTP {e.code}.{detail}", e.code) from e
except URLError as e:
- raise PlatformAPIError(f"Connection error: {e.reason}")
- except json.JSONDecodeError:
- raise PlatformAPIError("Invalid JSON response from platform")
+ raise PlatformAPIError(f"Connection error: {e.reason}") from e
+ except json.JSONDecodeError as e:
+ raise PlatformAPIError("Invalid JSON response from platform") from e
diff --git a/osmosis_ai/cli.py b/osmosis_ai/cli.py
index ef518722..23551e19 100644
--- a/osmosis_ai/cli.py
+++ b/osmosis_ai/cli.py
@@ -2,16 +2,22 @@
import argparse
import sys
-from typing import Optional
from dotenv import load_dotenv
-from .cli_commands import EvalRubricCommand, LoginCommand, LogoutCommand, PreviewCommand, WhoamiCommand, WorkspaceCommand
+from .cli_commands import (
+ EvalRubricCommand,
+ LoginCommand,
+ LogoutCommand,
+ PreviewCommand,
+ WhoamiCommand,
+ WorkspaceCommand,
+)
from .cli_services import CLIError
from .consts import PACKAGE_VERSION, package_name
-def main(argv: Optional[list[str]] = None) -> int:
+def main(argv: list[str] | None = None) -> int:
"""Entry point for the osmosis CLI."""
# Load environment variables from .env file in current working directory
load_dotenv()
@@ -36,14 +42,15 @@ def _build_parser() -> argparse.ArgumentParser:
prog="osmosis",
description="Osmosis AI SDK - rubric evaluation and remote rollout server.",
)
-
+
parser.add_argument(
- "-V", "--version",
+ "-V",
+ "--version",
action="version",
version=f"{package_name} {PACKAGE_VERSION}",
help="Show version number.",
)
-
+
subparsers = parser.add_subparsers(dest="command")
login_parser = subparsers.add_parser(
diff --git a/osmosis_ai/cli_commands.py b/osmosis_ai/cli_commands.py
index d946a4d9..7f85b447 100644
--- a/osmosis_ai/cli_commands.py
+++ b/osmosis_ai/cli_commands.py
@@ -2,28 +2,29 @@
import argparse
import time
+from collections.abc import Callable
from pathlib import Path
-from typing import Any, Callable, Optional
+from typing import Any
from .auth import (
LoginError,
delete_credentials,
delete_workspace_credentials,
- get_all_workspaces,
get_active_workspace,
+ get_all_workspaces,
load_credentials,
login,
set_active_workspace,
)
from .cli_services import (
- CLIError,
- ParsedItem,
BaselineComparator,
+ CLIError,
ConsoleReportRenderer,
DatasetLoader,
EvaluationSession,
EvaluationSessionRequest,
JsonReportWriter,
+ ParsedItem,
RubricEvaluationEngine,
RubricSuite,
discover_rubric_config_path,
@@ -74,7 +75,9 @@ def run(self, args: argparse.Namespace) -> int:
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.")
+ raise CLIError(
+ f"Unsupported file extension '{suffix}'. Expected .yaml, .yml, or .jsonl."
+ )
return 0
@@ -85,14 +88,16 @@ class EvalRubricCommand:
def __init__(
self,
*,
- session: Optional[EvaluationSession] = None,
- config_locator: Callable[[Optional[str], Path], Path] = discover_rubric_config_path,
+ session: EvaluationSession | None = None,
+ config_locator: Callable[
+ [str | None, 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,
+ dataset_loader: DatasetLoader | None = None,
+ engine: RubricEvaluationEngine | None = None,
+ renderer: ConsoleReportRenderer | None = None,
+ report_writer: JsonReportWriter | None = None,
+ baseline_comparator: BaselineComparator | None = None,
):
self._renderer = renderer or ConsoleReportRenderer()
if session is not None:
@@ -169,9 +174,15 @@ def run(self, args: argparse.Namespace) -> int:
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,
+ 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:
@@ -237,12 +248,18 @@ def run(self, args: argparse.Namespace) -> int:
print(f"\n[OK] Logged in as {result.user.email}")
if result.user.name:
print(f" Name: {result.user.name}")
- print(f" Workspace: {result.organization.name} ({result.organization.role})")
+ print(
+ f" Workspace: {result.organization.name} ({result.organization.role})"
+ )
print(f" Token expires: {result.expires_at.strftime('%Y-%m-%d')}")
if result.revoked_previous_tokens > 0:
- token_word = "token" if result.revoked_previous_tokens == 1 else "tokens"
- print(f" [Note] {result.revoked_previous_tokens} previous {token_word} for this device was revoked")
- print(f" Credentials saved to ~/.config/osmosis/credentials.json")
+ token_word = (
+ "token" if result.revoked_previous_tokens == 1 else "tokens"
+ )
+ print(
+ f" [Note] {result.revoked_previous_tokens} previous {token_word} for this device was revoked"
+ )
+ print(" Credentials saved to ~/.config/osmosis/credentials.json")
return 0
except LoginError as e:
@@ -284,7 +301,9 @@ def run(self, _args: argparse.Namespace) -> int:
for name, creds, is_active in workspaces:
active_marker = " *" if is_active else ""
expired_marker = " [expired]" if creds.is_expired() else ""
- print(f" {name} ({creds.organization.role}){active_marker}{expired_marker}")
+ print(
+ f" {name} ({creds.organization.role}){active_marker}{expired_marker}"
+ )
return 0
@@ -301,7 +320,8 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None:
help="Logout from all workspaces.",
)
parser.add_argument(
- "-y", "--yes",
+ "-y",
+ "--yes",
dest="skip_confirm",
action="store_true",
help="Skip confirmation prompt.",
@@ -369,7 +389,7 @@ def _logout_interactive(
expired_marker = " [expired]" if creds.is_expired() else ""
print(f" {i}. {name}{active_marker}{expired_marker}")
print(f" {len(workspaces) + 1}. All workspaces")
- print(f" 0. Cancel")
+ print(" 0. Cancel")
try:
choice = input("\nEnter number: ").strip()
@@ -402,19 +422,27 @@ class WorkspaceCommand:
"""Handler for `osmosis workspace`."""
def configure_parser(self, parser: argparse.ArgumentParser) -> None:
- subparsers = parser.add_subparsers(dest="workspace_action", help="Workspace management commands")
+ subparsers = parser.add_subparsers(
+ dest="workspace_action", help="Workspace management commands"
+ )
# workspace list
- list_parser = subparsers.add_parser("list", help="List all logged-in workspaces")
+ list_parser = subparsers.add_parser(
+ "list", help="List all logged-in workspaces"
+ )
list_parser.set_defaults(handler=self._run_list)
# workspace switch
- switch_parser = subparsers.add_parser("switch", help="Switch to a different workspace")
+ switch_parser = subparsers.add_parser(
+ "switch", help="Switch to a different workspace"
+ )
switch_parser.add_argument("name", help="Name of the workspace to switch to")
switch_parser.set_defaults(handler=self._run_switch)
# workspace current
- current_parser = subparsers.add_parser("current", help="Show the current active workspace")
+ current_parser = subparsers.add_parser(
+ "current", help="Show the current active workspace"
+ )
current_parser.set_defaults(handler=self._run_current)
# Default handler when no subcommand is provided
@@ -435,14 +463,18 @@ def _run_list(self, args: argparse.Namespace) -> int:
workspaces = get_all_workspaces()
if not workspaces:
- print("No workspaces logged in. Run 'osmosis login' to log in to a workspace.")
+ print(
+ "No workspaces logged in. Run 'osmosis login' to log in to a workspace."
+ )
return 0
print("Logged-in workspaces:")
for name, creds, is_active in workspaces:
active_marker = " (active)" if is_active else ""
expired_marker = " [expired]" if creds.is_expired() else ""
- print(f" {name} ({creds.organization.role}){active_marker}{expired_marker}")
+ print(
+ f" {name} ({creds.organization.role}){active_marker}{expired_marker}"
+ )
return 0
@@ -470,7 +502,9 @@ def _run_current(self, args: argparse.Namespace) -> int:
for name, creds, is_active in workspaces:
if is_active:
expired_marker = " [expired]" if creds.is_expired() else ""
- print(f"Current workspace: {name} ({creds.organization.role}){expired_marker}")
+ print(
+ f"Current workspace: {name} ({creds.organization.role}){expired_marker}"
+ )
print(f" User: {creds.user.email}")
print(f" Expires: {creds.expires_at.strftime('%Y-%m-%d')}")
return 0
diff --git a/osmosis_ai/cli_services/__init__.py b/osmosis_ai/cli_services/__init__.py
index 9e13279f..61a4f7c9 100644
--- a/osmosis_ai/cli_services/__init__.py
+++ b/osmosis_ai/cli_services/__init__.py
@@ -10,7 +10,12 @@
load_rubric_suite,
render_yaml_items,
)
-from .dataset import DatasetLoader, DatasetRecord, load_jsonl_records, render_json_records
+from .dataset import (
+ DatasetLoader,
+ DatasetRecord,
+ load_jsonl_records,
+ render_json_records,
+)
from .engine import (
EvaluationRecordResult,
EvaluationReport,
@@ -27,7 +32,11 @@
JsonReportWriter,
TextReportFormatter,
)
-from .session import EvaluationSession, EvaluationSessionRequest, EvaluationSessionResult
+from .session import (
+ EvaluationSession,
+ EvaluationSessionRequest,
+ EvaluationSessionResult,
+)
__all__ = [
"BaselineComparator",
@@ -36,12 +45,12 @@
"ConsoleReportRenderer",
"DatasetLoader",
"DatasetRecord",
- "EvaluationSession",
- "EvaluationSessionRequest",
- "EvaluationSessionResult",
"EvaluationRecordResult",
"EvaluationReport",
"EvaluationRun",
+ "EvaluationSession",
+ "EvaluationSessionRequest",
+ "EvaluationSessionResult",
"JsonReportFormatter",
"JsonReportWriter",
"ParsedItem",
diff --git a/osmosis_ai/cli_services/config.py b/osmosis_ai/cli_services/config.py
index d26094c1..c94ea819 100644
--- a/osmosis_ai/cli_services/config.py
+++ b/osmosis_ai/cli_services/config.py
@@ -1,9 +1,10 @@
from __future__ import annotations
import copy
+from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
-from typing import Any, Optional, Sequence
+from typing import Any
import yaml
from yaml.representer import SafeRepresenter
@@ -14,7 +15,7 @@
@dataclass(frozen=True)
class ParsedItem:
- label: Optional[str]
+ label: str | None
payload: Any
@@ -23,18 +24,18 @@ class RubricConfig:
rubric_id: str
rubric_text: str
model_info: dict[str, Any]
- score_min: Optional[float]
- score_max: Optional[float]
- system_prompt: Optional[str]
- original_input: Optional[str]
- ground_truth: Optional[str]
+ score_min: float | None
+ score_max: float | None
+ system_prompt: str | None
+ original_input: str | None
+ ground_truth: str | None
source_label: str
@dataclass(frozen=True)
class RubricSuite:
source_path: Path
- version: Optional[int]
+ version: int | None
configs: dict[str, RubricConfig]
def get(self, rubric_id: str) -> RubricConfig:
@@ -58,7 +59,7 @@ class RubricConfigDocumentResult:
class RubricConfigDocumentSchema:
"""Base interface for schema-specific rubric config parsing."""
- version: Optional[int] = None
+ version: int | None = None
def parse_document(
self,
@@ -86,7 +87,9 @@ def parse_document(
) -> 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)
+ return _build_document_configs(
+ entries, defaults, path=path, doc_index=doc_index, strict=strict
+ )
class Version1RubricConfigSchema(BaseRubricConfigSchema):
@@ -98,7 +101,11 @@ class Version1RubricConfigSchema(BaseRubricConfigSchema):
class RubricConfigParser:
"""Parses rubric configuration files and produces typed suites."""
- def __init__(self, *, schemas: Optional[dict[Optional[int], RubricConfigDocumentSchema]] = None):
+ def __init__(
+ self,
+ *,
+ schemas: dict[int | None, RubricConfigDocumentSchema] | None = None,
+ ):
self._schemas = schemas or {
None: BaseRubricConfigSchema(),
1: Version1RubricConfigSchema(),
@@ -106,11 +113,13 @@ def __init__(self, *, schemas: Optional[dict[Optional[int], RubricConfigDocument
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]]:
+ 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
+ detected_version: int | None = None
document_indices = []
for doc_index, document in enumerate(documents):
@@ -144,7 +153,9 @@ def parse(self, path: Path, *, strict: bool = True) -> tuple[RubricSuite, list[P
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}'.")
+ raise CLIError(
+ f"Duplicate rubric id '{rubric_id}' detected in '{path}'."
+ )
configs[rubric_id] = config
if strict and not configs:
@@ -153,7 +164,7 @@ def parse(self, path: Path, *, strict: bool = True) -> tuple[RubricSuite, list[P
suite = RubricSuite(source_path=path, version=detected_version, configs=configs)
return suite, parsed_items
- def _select_schema(self, version: Optional[int]) -> RubricConfigDocumentSchema:
+ def _select_schema(self, version: int | None) -> RubricConfigDocumentSchema:
if version in self._schemas:
return self._schemas[version]
if version is None:
@@ -161,7 +172,9 @@ def _select_schema(self, version: Optional[int]) -> RubricConfigDocumentSchema:
raise CLIError(f"Unsupported rubric config version '{version}'.")
@staticmethod
- def _coerce_optional_version(document: Any, path: Path, doc_index: int) -> Optional[int]:
+ def _coerce_optional_version(
+ document: Any, path: Path, doc_index: int
+ ) -> int | None:
if not isinstance(document, dict):
return None
version_value = document.get("version")
@@ -195,9 +208,7 @@ def _build_document_configs(
if not isinstance(payload, dict):
continue
if "extra_info" in payload:
- message = (
- f"Rubric entry in '{path}' (document {doc_index + 1}) must not include 'extra_info'."
- )
+ message = f"Rubric entry in '{path}' (document {doc_index + 1}) must not include 'extra_info'."
if strict:
raise CLIError(message)
continue
@@ -271,7 +282,7 @@ def _build_document_configs(
return RubricConfigDocumentResult(configs=configs, items=parsed_items)
-def discover_rubric_config_path(config_arg: Optional[str], data_path: Path) -> Path:
+def discover_rubric_config_path(config_arg: str | None, data_path: Path) -> Path:
if config_arg:
candidate = Path(config_arg).expanduser()
if not candidate.exists():
@@ -344,7 +355,9 @@ def _load_yaml_documents(path: Path) -> list[Any]:
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]:
+def _extract_config_defaults(
+ document: Any, path: Path, doc_index: int
+) -> dict[str, Any]:
if not isinstance(document, dict):
return {
"model_info": None,
@@ -375,7 +388,9 @@ def _extract_config_defaults(document: Any, path: Path, doc_index: int) -> dict[
return defaults
-def _extract_rubric_items(node: Any, context: Optional[str], doc_index: int) -> list[ParsedItem]:
+def _extract_rubric_items(
+ node: Any, context: str | None, doc_index: int
+) -> list[ParsedItem]:
items: list[ParsedItem] = []
if node is None:
@@ -388,11 +403,17 @@ def _extract_rubric_items(node: Any, context: Optional[str], doc_index: int) ->
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))
+ 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))
+ items.extend(
+ _extract_rubric_items(value, context=idx_context, doc_index=doc_index)
+ )
return items
diff --git a/osmosis_ai/cli_services/dataset.py b/osmosis_ai/cli_services/dataset.py
index 1e0119bd..fd8ec0b6 100644
--- a/osmosis_ai/cli_services/dataset.py
+++ b/osmosis_ai/cli_services/dataset.py
@@ -2,9 +2,10 @@
import copy
import json
+from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
-from typing import Any, Optional, Sequence
+from typing import Any
from .errors import CLIError
from .shared import coerce_optional_float
@@ -14,17 +15,17 @@
class DatasetRecord:
payload: dict[str, Any]
rubric_id: str
- conversation_id: Optional[str]
- record_id: Optional[str]
+ conversation_id: str | None
+ record_id: str | None
solution_str: str
- ground_truth: 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 merged_extra_info(self) -> Optional[dict[str, Any]]:
+ ground_truth: str | None
+ original_input: str | None
+ metadata: dict[str, Any] | None
+ extra_info: dict[str, Any] | None
+ score_min: float | None
+ score_max: float | None
+
+ def merged_extra_info(self) -> dict[str, Any] | None:
merged: dict[str, Any] = {}
if isinstance(self.extra_info, dict):
merged.update(copy.deepcopy(self.extra_info))
@@ -32,7 +33,7 @@ def merged_extra_info(self) -> Optional[dict[str, Any]]:
merged.setdefault("dataset_metadata", copy.deepcopy(self.metadata))
return merged or None
- def assistant_preview(self, *, max_length: int = 140) -> Optional[str]:
+ def assistant_preview(self, *, max_length: int = 140) -> str | None:
text = self.solution_str.strip()
if not text:
return None
@@ -98,21 +99,37 @@ def _create_record(payload: dict[str, Any]) -> DatasetRecord:
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
+ 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 ''}'"
+ payload.get("score_min"),
+ "score_min",
+ f"record '{conversation_id or rubric_id or ''}'",
)
score_max = coerce_optional_float(
- payload.get("score_max"), "score_max", f"record '{conversation_id or rubric_id or ''}'"
+ payload.get("score_max"),
+ "score_max",
+ f"record '{conversation_id or rubric_id or ''}'",
)
- 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
+ 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 ""
solution_raw = payload.get("solution_str")
if not isinstance(solution_raw, str) or not solution_raw.strip():
- raise CLIError(f"Record '{record_label}' must include a non-empty 'solution_str' string.")
+ raise CLIError(
+ f"Record '{record_label}' must include a non-empty 'solution_str' string."
+ )
original_input_raw = payload.get("original_input")
if isinstance(original_input_raw, str):
@@ -131,7 +148,9 @@ def _create_record(payload: dict[str, Any]) -> DatasetRecord:
conversation_id=conversation_id,
record_id=record_id,
solution_str=solution_raw,
- ground_truth=payload.get("ground_truth") if isinstance(payload.get("ground_truth"), str) else None,
+ ground_truth=payload.get("ground_truth")
+ if isinstance(payload.get("ground_truth"), str)
+ else None,
original_input=original_input,
metadata=metadata,
extra_info=extra_info,
@@ -150,9 +169,13 @@ def load_jsonl_records(path: Path) -> list[dict[str, Any]]:
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
+ 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}'.")
+ raise CLIError(
+ f"Expected JSON object on line {line_number} of '{path}'."
+ )
records.append(record)
if not records:
diff --git a/osmosis_ai/cli_services/engine.py b/osmosis_ai/cli_services/engine.py
index 2d839f20..9481b36b 100644
--- a/osmosis_ai/cli_services/engine.py
+++ b/osmosis_ai/cli_services/engine.py
@@ -1,13 +1,14 @@
from __future__ import annotations
import copy
+import inspect
import sys
import time
+from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
-import inspect
-from typing import Any, Optional, Sequence
+from typing import Any
from tqdm import tqdm
@@ -19,7 +20,7 @@
from .shared import calculate_statistics, coerce_optional_float, collapse_preview_text
-def _normalize_config_str(value: Any) -> Optional[str]:
+def _normalize_config_str(value: Any) -> str | None:
if value is None:
return None
text = str(value).strip()
@@ -27,19 +28,19 @@ def _normalize_config_str(value: Any) -> Optional[str]:
def _compose_extra_info_context(
- base: Optional[dict[str, Any]],
+ base: dict[str, Any] | None,
*,
rubric_text: str,
- provider: Optional[str],
- model: Optional[str],
- system_prompt: Optional[str],
- original_input: Optional[str],
- api_key: Optional[str],
- api_key_env: Optional[str],
- score_min: Optional[float],
- score_max: Optional[float],
+ provider: str | None,
+ model: str | None,
+ system_prompt: str | None,
+ original_input: str | None,
+ api_key: str | None,
+ api_key_env: str | None,
+ score_min: float | None,
+ score_max: float | None,
model_info: dict[str, Any],
-) -> tuple[dict[str, Any], Optional[dict[str, Any]]]:
+) -> tuple[dict[str, Any], dict[str, Any] | None]:
"""
Build the runtime context passed to rubric functions along with a sanitised
copy safe for prompt injection.
@@ -76,7 +77,7 @@ def _compose_extra_info_context(
model_info_copy["api_key_env"] = api_key_env
decorated_payload["model_info"] = model_info_copy
- prompt_payload: Optional[dict[str, Any]] = None
+ prompt_payload: dict[str, Any] | None = None
base_metadata = decorated_payload.get("metadata")
if isinstance(base_metadata, dict):
prompt_payload = copy.deepcopy(base_metadata)
@@ -98,9 +99,9 @@ def _compose_extra_info_context(
def _merge_system_prompts(
- prepend_prompt: Optional[str],
- base_prompt: Optional[str],
-) -> Optional[str]:
+ prepend_prompt: str | None,
+ base_prompt: str | None,
+) -> str | None:
prompts: list[str] = []
if prepend_prompt:
prompts.append(prepend_prompt)
@@ -121,7 +122,9 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
solution = record.solution_str
if not isinstance(solution, str) or not solution.strip():
label = record.conversation_id or record.rubric_id or ""
- raise CLIError(f"Record '{label}' must include a non-empty 'solution_str' string.")
+ raise CLIError(
+ f"Record '{label}' must include a non-empty 'solution_str' string."
+ )
score_min = coerce_optional_float(
record.score_min if record.score_min is not None else config.score_min,
@@ -134,8 +137,16 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
f"record '{record.conversation_id or ''}'",
)
- ground_truth = record.ground_truth if record.ground_truth is not None else config.ground_truth
- original_input = record.original_input if record.original_input is not None else config.original_input
+ ground_truth = (
+ record.ground_truth
+ if record.ground_truth is not None
+ else config.ground_truth
+ )
+ original_input = (
+ record.original_input
+ if record.original_input is not None
+ else config.original_input
+ )
provider_value = _normalize_config_str(config.model_info.get("provider"))
model_value = _normalize_config_str(config.model_info.get("model"))
@@ -149,8 +160,12 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
try:
model_info_payload = copy.deepcopy(config.model_info)
- base_system_prompt = _normalize_config_str(model_info_payload.get("system_prompt"))
- combined_system_prompt = _merge_system_prompts(system_prompt_value, base_system_prompt)
+ base_system_prompt = _normalize_config_str(
+ model_info_payload.get("system_prompt")
+ )
+ combined_system_prompt = _merge_system_prompts(
+ system_prompt_value, base_system_prompt
+ )
if combined_system_prompt is not None:
model_info_payload["system_prompt"] = combined_system_prompt
else:
@@ -173,9 +188,14 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
signature = inspect.signature(self._evaluate_fn)
parameters = signature.parameters
accepts_var_kwargs = any(
- param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()
+ param.kind == inspect.Parameter.VAR_KEYWORD
+ for param in parameters.values()
+ )
+ is_evaluate_rubric_style = (
+ accepts_var_kwargs
+ or "rubric" in parameters
+ or "model_info" in parameters
)
- is_evaluate_rubric_style = accepts_var_kwargs or "rubric" in parameters or "model_info" in parameters
if is_evaluate_rubric_style:
return self._evaluate_fn(
@@ -193,7 +213,10 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
call_args: list[Any] = []
call_kwargs: dict[str, Any] = {}
for param in parameters.values():
- if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+ if param.kind in (
+ inspect.Parameter.VAR_POSITIONAL,
+ inspect.Parameter.VAR_KEYWORD,
+ ):
continue
if param.name == "solution_str":
@@ -222,13 +245,13 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
class EvaluationRun:
run_index: int
status: str
- score: Optional[float]
- explanation: Optional[str]
- preview: Optional[str]
+ score: float | None
+ explanation: str | None
+ preview: str | None
duration_seconds: float
started_at: datetime
completed_at: datetime
- error: Optional[str]
+ error: str | None
raw: Any
@@ -254,7 +277,7 @@ class EvaluationReport:
class RubricEvaluationEngine:
"""Executes rubric evaluations across a dataset and aggregates statistics."""
- def __init__(self, evaluator: Optional[RubricEvaluator] = None):
+ def __init__(self, evaluator: RubricEvaluator | None = None):
self._evaluator = evaluator or RubricEvaluator()
def execute(
@@ -272,7 +295,9 @@ def execute(
total_successes = 0
progress_total = len(records) * number
- show_progress = progress_total > 1 and getattr(sys.stderr, "isatty", lambda: False)()
+ show_progress = (
+ progress_total > 1 and getattr(sys.stderr, "isatty", lambda: False)()
+ )
progress = (
tqdm(
total=progress_total,
@@ -296,10 +321,10 @@ def execute(
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
+ error_message: str | None = None
+ score_value: float | None = None
+ explanation_value: str | None = None
+ preview_value: str | None = None
raw_payload: Any = None
try:
@@ -320,8 +345,12 @@ def execute(
if 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)
+ 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)
@@ -389,7 +418,9 @@ def execute(
)
@staticmethod
- def _resolve_preview_text(result: Optional[dict[str, Any]], fallback: Optional[str]) -> Optional[str]:
+ def _resolve_preview_text(
+ result: dict[str, Any] | None, fallback: str | None
+ ) -> str | None:
if not isinstance(result, dict):
return fallback
preview = collapse_preview_text(result.get("preview"))
@@ -405,7 +436,7 @@ def _resolve_preview_text(result: Optional[dict[str, Any]], fallback: Optional[s
return fallback
-def _extract_float(value: Any) -> Optional[float]:
+def _extract_float(value: Any) -> float | None:
try:
if isinstance(value, (int, float)) and not isinstance(value, bool):
return float(value)
@@ -414,7 +445,7 @@ def _extract_float(value: Any) -> Optional[float]:
return None
-def _normalize_optional_text(value: Any) -> Optional[str]:
+def _normalize_optional_text(value: Any) -> str | None:
if value is None:
return None
text = str(value).strip()
diff --git a/osmosis_ai/cli_services/reporting.py b/osmosis_ai/cli_services/reporting.py
index 8497b534..b448513d 100644
--- a/osmosis_ai/cli_services/reporting.py
+++ b/osmosis_ai/cli_services/reporting.py
@@ -1,9 +1,10 @@
from __future__ import annotations
import json
+from collections.abc import Callable
from datetime import datetime, timezone
from pathlib import Path
-from typing import Any, Callable, Optional
+from typing import Any
from .engine import EvaluationRecordResult, EvaluationReport, EvaluationRun
from .errors import CLIError
@@ -16,17 +17,25 @@ class TextReportFormatter:
def build(
self,
report: EvaluationReport,
- baseline: Optional["BaselineStatistics"] = None,
+ baseline: BaselineStatistics | None = None,
) -> list[str]:
lines: list[str] = []
- provider = str(report.rubric_config.model_info.get("provider", "")).strip() or ""
- model_name = str(report.rubric_config.model_info.get("model", "")).strip() or ""
+ provider = (
+ str(report.rubric_config.model_info.get("provider", "")).strip()
+ or ""
+ )
+ model_name = (
+ str(report.rubric_config.model_info.get("model", "")).strip()
+ or ""
+ )
lines.append(
f"Rubric '{report.rubric_config.rubric_id}' "
f"({report.rubric_config.source_label}) -> provider '{provider}' model '{model_name}'"
)
- lines.append(f"Loaded {len(report.record_results)} matching record(s) from {report.data_path}")
+ lines.append(
+ f"Loaded {len(report.record_results)} matching record(s) from {report.data_path}"
+ )
lines.append(f"Running {report.number} evaluation(s) per record")
lines.append("")
@@ -44,7 +53,9 @@ def _format_record(self, record_result: EvaluationRecordResult) -> list[str]:
if index < total_runs - 1:
lines.append("")
- summary_lines = self._format_summary(record_result.statistics, len(record_result.runs))
+ summary_lines = self._format_summary(
+ record_result.statistics, len(record_result.runs)
+ )
if summary_lines:
lines.extend(summary_lines)
lines.append("")
@@ -66,26 +77,36 @@ def _format_run(self, run: EvaluationRun) -> list[str]:
lines.append(self._format_detail_line("preview", run.preview))
if run.explanation:
lines.append(self._format_detail_line("explanation", run.explanation))
- lines.append(self._format_detail_line("duration", f"{run.duration_seconds:.2f}s"))
+ lines.append(
+ self._format_detail_line("duration", f"{run.duration_seconds:.2f}s")
+ )
return lines
- def _format_summary(self, statistics: dict[str, float], total_runs: int) -> list[str]:
+ def _format_summary(
+ self, statistics: dict[str, float], total_runs: int
+ ) -> list[str]:
if not statistics:
return []
- success_count = int(round(statistics.get("success_count", total_runs)))
- failure_count = int(round(statistics.get("failure_count", total_runs - success_count)))
+ success_count = round(statistics.get("success_count", total_runs))
+ failure_count = round(
+ statistics.get("failure_count", total_runs - success_count)
+ )
if total_runs <= 1 and failure_count == 0:
return []
lines = [" Summary:"]
- lines.append(f" total: {int(round(statistics.get('total_runs', total_runs)))}")
+ lines.append(
+ f" total: {round(statistics.get('total_runs', total_runs))}"
+ )
lines.append(f" successes: {success_count}")
lines.append(f" failures: {failure_count}")
if success_count > 0:
lines.append(f" average: {statistics.get('average', 0.0):.4f}")
lines.append(f" variance: {statistics.get('variance', 0.0):.6f}")
lines.append(f" stdev: {statistics.get('stdev', 0.0):.4f}")
- lines.append(f" min/max: {statistics.get('min', 0.0):.4f} / {statistics.get('max', 0.0):.4f}")
+ lines.append(
+ f" min/max: {statistics.get('min', 0.0):.4f} / {statistics.get('max', 0.0):.4f}"
+ )
else:
lines.append(" average: n/a")
lines.append(" variance: n/a")
@@ -96,11 +117,20 @@ def _format_summary(self, statistics: dict[str, float], total_runs: int) -> list
def _format_baseline(
self,
report: EvaluationReport,
- baseline: "BaselineStatistics",
+ baseline: BaselineStatistics,
) -> list[str]:
lines = [f"Baseline comparison (source: {baseline.source_path}):"]
deltas = baseline.delta(report.overall_statistics)
- keys = ["average", "variance", "stdev", "min", "max", "success_count", "failure_count", "total_runs"]
+ keys = [
+ "average",
+ "variance",
+ "stdev",
+ "min",
+ "max",
+ "success_count",
+ "failure_count",
+ "total_runs",
+ ]
for key in keys:
if key not in baseline.statistics or key not in report.overall_statistics:
@@ -110,8 +140,8 @@ def _format_baseline(
delta_value = float(deltas.get(key, current_value - baseline_value))
if key in {"success_count", "failure_count", "total_runs"}:
- baseline_str = f"{int(round(baseline_value))}"
- current_str = f"{int(round(current_value))}"
+ baseline_str = f"{round(baseline_value)}"
+ current_str = f"{round(current_value)}"
delta_str = f"{delta_value:+.0f}"
else:
precision = 6 if key == "variance" else 4
@@ -139,7 +169,7 @@ class ConsoleReportRenderer:
def __init__(
self,
printer: Callable[[str], None] = print,
- formatter: Optional[TextReportFormatter] = None,
+ formatter: TextReportFormatter | None = None,
):
self._printer = printer
self._formatter = formatter or TextReportFormatter()
@@ -147,7 +177,7 @@ def __init__(
def render(
self,
report: EvaluationReport,
- baseline: Optional["BaselineStatistics"] = None,
+ baseline: BaselineStatistics | None = None,
) -> None:
for line in self._formatter.build(report, baseline):
self._printer(line)
@@ -169,7 +199,9 @@ def load(self, path: Path) -> BaselineStatistics:
if not path.exists():
raise CLIError(f"Baseline path '{path}' does not exist.")
if path.is_dir():
- raise CLIError(f"Baseline path '{path}' is a directory, expected JSON file.")
+ raise CLIError(
+ f"Baseline path '{path}' is a directory, expected JSON file."
+ )
try:
payload = json.loads(path.read_text(encoding="utf-8"))
@@ -196,7 +228,9 @@ def load(self, path: Path) -> BaselineStatistics:
except (TypeError, ValueError):
continue
if not statistics:
- raise CLIError("Baseline statistics could not be parsed into numeric values.")
+ raise CLIError(
+ "Baseline statistics could not be parsed into numeric values."
+ )
return BaselineStatistics(source_path=path, statistics=statistics)
@@ -208,11 +242,17 @@ def build(
self,
report: EvaluationReport,
*,
- output_identifier: Optional[str],
- baseline: Optional[BaselineStatistics],
+ output_identifier: str | None,
+ baseline: BaselineStatistics | None,
) -> dict[str, Any]:
- provider = str(report.rubric_config.model_info.get("provider", "")).strip() or ""
- model_name = str(report.rubric_config.model_info.get("model", "")).strip() or ""
+ provider = (
+ str(report.rubric_config.model_info.get("provider", "")).strip()
+ or ""
+ )
+ model_name = (
+ str(report.rubric_config.model_info.get("model", "")).strip()
+ or ""
+ )
generated: dict[str, Any] = {
"generated_at": datetime.now(timezone.utc).isoformat(),
@@ -231,7 +271,9 @@ def build(
for record_result in report.record_results:
conversation_label = record_result.conversation_label
- record_identifier = record_result.record.record_identifier(conversation_label)
+ record_identifier = record_result.record.record_identifier(
+ conversation_label
+ )
record_payload: dict[str, Any] = {
"id": record_identifier,
@@ -264,7 +306,9 @@ def build(
generated["baseline_comparison"] = {
"source_path": str(baseline.source_path),
"baseline_statistics": _normalise_statistics(baseline.statistics),
- "delta_statistics": _normalise_statistics(baseline.delta(report.overall_statistics)),
+ "delta_statistics": _normalise_statistics(
+ baseline.delta(report.overall_statistics)
+ ),
}
return generated
@@ -273,7 +317,7 @@ def build(
class JsonReportWriter:
"""Serialises an evaluation report to disk."""
- def __init__(self, formatter: Optional[JsonReportFormatter] = None):
+ def __init__(self, formatter: JsonReportFormatter | None = None):
self._formatter = formatter or JsonReportFormatter()
def write(
@@ -281,8 +325,8 @@ def write(
report: EvaluationReport,
*,
output_path: Path,
- output_identifier: Optional[str],
- baseline: Optional[BaselineStatistics],
+ output_identifier: str | None,
+ baseline: BaselineStatistics | None,
) -> Path:
parent_dir = output_path.parent
if parent_dir and not parent_dir.exists():
@@ -297,11 +341,12 @@ def write(
json.dump(payload, fh, indent=2, ensure_ascii=False)
return output_path
+
def _normalise_statistics(stats: dict[str, float]) -> dict[str, Any]:
normalised: dict[str, Any] = {}
for key, value in stats.items():
if key in {"success_count", "failure_count", "total_runs"}:
- normalised[key] = int(round(value))
+ normalised[key] = round(value)
else:
normalised[key] = float(value)
return normalised
diff --git a/osmosis_ai/cli_services/session.py b/osmosis_ai/cli_services/session.py
index a90f90b5..f21103e3 100644
--- a/osmosis_ai/cli_services/session.py
+++ b/osmosis_ai/cli_services/session.py
@@ -1,25 +1,31 @@
from __future__ import annotations
import time
+from collections.abc import Callable, Sequence
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 .config import (
+ RubricConfig,
+ RubricSuite,
+ discover_rubric_config_path,
+ load_rubric_suite,
+)
from .dataset import DatasetLoader, DatasetRecord
-from .engine import RubricEvaluationEngine, EvaluationReport
+from .engine import EvaluationReport, RubricEvaluationEngine
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 = "".join(
+ ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in rubric_id.strip()
+ )
clean = clean.strip("_") or "rubric"
return clean.lower()
@@ -29,10 +35,10 @@ 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
+ config_path: Path | None = None
+ output_path: Path | None = None
+ output_identifier: str | None = None
+ baseline_path: Path | None = None
@dataclass
@@ -43,9 +49,9 @@ class EvaluationSessionResult:
rubric_config: RubricConfig
records: Sequence[DatasetRecord]
report: EvaluationReport
- baseline: Optional[BaselineStatistics]
- written_path: Optional[Path]
- output_identifier: Optional[str]
+ baseline: BaselineStatistics | None
+ written_path: Path | None
+ output_identifier: str | None
class EvaluationSession:
@@ -54,13 +60,15 @@ class EvaluationSession:
def __init__(
self,
*,
- config_locator: Callable[[Optional[str], Path], Path] = discover_rubric_config_path,
+ config_locator: Callable[
+ [str | None, 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,
+ dataset_loader: DatasetLoader | None = None,
+ engine: RubricEvaluationEngine | None = None,
+ baseline_comparator: BaselineComparator | None = None,
+ report_writer: JsonReportWriter | None = None,
+ identifier_factory: Callable[[], str] | None = None,
):
self._config_locator = config_locator
self._suite_loader = suite_loader
@@ -84,9 +92,13 @@ def execute(self, request: EvaluationSessionRequest) -> EvaluationSessionResult:
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}'.")
+ 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_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)
@@ -98,10 +110,14 @@ def execute(self, request: EvaluationSessionRequest) -> EvaluationSessionResult:
all_records = self._dataset_loader.load(data_path)
matching_records = [
- record for record in all_records if record.rubric_id.lower() == rubric_id.lower()
+ 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}'.")
+ raise CLIError(
+ f"No records in '{data_path}' reference rubric '{rubric_id}'."
+ )
baseline_stats = self._load_baseline(request.baseline_path)
@@ -140,7 +156,7 @@ def execute(self, request: EvaluationSessionRequest) -> EvaluationSessionResult:
output_identifier=resolved_identifier,
)
- def _load_baseline(self, baseline_path: Optional[Path]) -> Optional[BaselineStatistics]:
+ def _load_baseline(self, baseline_path: Path | None) -> BaselineStatistics | None:
if baseline_path is None:
return None
resolved = baseline_path.expanduser()
@@ -148,11 +164,11 @@ def _load_baseline(self, baseline_path: Optional[Path]) -> Optional[BaselineStat
def _resolve_output_path(
self,
- output_candidate: Optional[Path],
- output_identifier: Optional[str],
+ output_candidate: Path | None,
+ output_identifier: str | None,
*,
rubric_id: str,
- ) -> tuple[Optional[Path], Optional[str]]:
+ ) -> tuple[Path | None, str | None]:
if output_candidate is None:
identifier = output_identifier or self._identifier_factory()
target_dir = _CACHE_ROOT / _sanitise_rubric_folder(rubric_id)
diff --git a/osmosis_ai/cli_services/shared.py b/osmosis_ai/cli_services/shared.py
index 14a15b4f..2f6ea730 100644
--- a/osmosis_ai/cli_services/shared.py
+++ b/osmosis_ai/cli_services/shared.py
@@ -1,12 +1,15 @@
from __future__ import annotations
-from statistics import mean, pvariance, pstdev
-from typing import Any, Collection, Optional, Set
+from collections.abc import Collection
+from statistics import mean, pstdev, pvariance
+from typing import Any
from .errors import CLIError
-def coerce_optional_float(value: Any, field_name: str, source_label: str) -> Optional[float]:
+def coerce_optional_float(
+ value: Any, field_name: str, source_label: str
+) -> float | None:
if value is None:
return None
if isinstance(value, (int, float)) and not isinstance(value, bool):
@@ -16,7 +19,7 @@ def coerce_optional_float(value: Any, field_name: str, source_label: str) -> Opt
)
-def collapse_preview_text(value: Any, *, max_length: int = 140) -> Optional[str]:
+def collapse_preview_text(value: Any, *, max_length: int = 140) -> str | None:
if not isinstance(value, str):
return None
collapsed = " ".join(value.strip().split())
@@ -48,7 +51,9 @@ def calculate_statistics(scores: list[float]) -> dict[str, float]:
}
-def calculate_stat_deltas(baseline: dict[str, float], current: dict[str, float]) -> dict[str, float]:
+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:
@@ -67,8 +72,8 @@ def gather_text_fragments(
fragments: list[str],
*,
allow_free_strings: bool = False,
- seen: Optional[Set[int]] = None,
- string_key_allowlist: Optional[Collection[str]] = None,
+ seen: set[int] | None = None,
+ string_key_allowlist: Collection[str] | None = None,
) -> None:
"""Collect textual snippets from nested message-like structures.
@@ -115,7 +120,7 @@ def gather_text_fragments(
allowlist = {key.lower() for key in allowlist}
prioritized_keys = ("text", "value")
- handled_keys: Set[str] = {
+ handled_keys: set[str] = {
"text",
"value",
"content",
@@ -148,15 +153,7 @@ def gather_text_fragments(
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:
+ if (node.get("type") == "tool_result" and "content" in node) or "content" in node:
gather_text_fragments(
node["content"],
fragments,
@@ -196,7 +193,7 @@ def collect_text_fragments(
node: Any,
*,
allow_free_strings: bool = False,
- string_key_allowlist: Optional[Collection[str]] = None,
+ string_key_allowlist: Collection[str] | None = None,
) -> list[str]:
fragments: list[str] = []
gather_text_fragments(
diff --git a/osmosis_ai/rollout/__init__.py b/osmosis_ai/rollout/__init__.py
index 4f4d92fb..515d58b0 100644
--- a/osmosis_ai/rollout/__init__.py
+++ b/osmosis_ai/rollout/__init__.py
@@ -39,16 +39,25 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
"""
# Core classes
+from osmosis_ai.rollout.client import (
+ CompletionsResult,
+ OsmosisLLMClient,
+)
+
+# Configuration
+from osmosis_ai.rollout.config import (
+ RolloutClientSettings,
+ RolloutServerSettings,
+ RolloutSettings,
+ configure,
+ get_settings,
+ reset_settings,
+)
from osmosis_ai.rollout.core.base import (
RolloutAgentLoop,
RolloutContext,
RolloutResult,
)
-from osmosis_ai.rollout.client import (
- CompletionsResult,
- OsmosisLLMClient,
-)
-from osmosis_ai.rollout.core.llm_client import LLMClientProtocol
from osmosis_ai.rollout.core.exceptions import (
AgentLoopNotFoundError,
OsmosisRolloutError,
@@ -59,19 +68,13 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
ToolArgumentError,
ToolExecutionError,
)
-from osmosis_ai.rollout.registry import (
- AgentLoopRegistry,
- get_agent_loop,
- list_agent_loops,
- register_agent_loop,
- unregister_agent_loop,
-)
+from osmosis_ai.rollout.core.llm_client import LLMClientProtocol
from osmosis_ai.rollout.core.schemas import (
- CompletionUsage,
+ DEFAULT_MAX_METADATA_SIZE_BYTES,
CompletionsChoice,
CompletionsRequest,
CompletionsResponse,
- DEFAULT_MAX_METADATA_SIZE_BYTES,
+ CompletionUsage,
InitResponse,
MessageDict,
OpenAIFunctionCallSchema,
@@ -90,6 +93,20 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
get_max_metadata_size_bytes,
set_max_metadata_size_bytes,
)
+from osmosis_ai.rollout.network import (
+ PublicIPDetectionError,
+ detect_public_ip,
+ is_private_ip,
+ is_valid_hostname_or_ip,
+ validate_ipv4,
+)
+from osmosis_ai.rollout.registry import (
+ AgentLoopRegistry,
+ get_agent_loop,
+ list_agent_loops,
+ register_agent_loop,
+ unregister_agent_loop,
+)
from osmosis_ai.rollout.server.app import create_app
from osmosis_ai.rollout.server.serve import (
DEFAULT_HOST,
@@ -98,12 +115,6 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
serve_agent_loop,
validate_and_report,
)
-from osmosis_ai.rollout.validator import (
- AgentLoopValidationError,
- ValidationError,
- ValidationResult,
- validate_agent_loop,
-)
from osmosis_ai.rollout.tools import (
create_tool_error_result,
create_tool_result,
@@ -122,115 +133,104 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
normalize_stop,
parse_tool_calls,
)
-from osmosis_ai.rollout.network import (
- detect_public_ip,
- PublicIPDetectionError,
- validate_ipv4,
- is_valid_hostname_or_ip,
- is_private_ip,
-)
-
-# Configuration
-from osmosis_ai.rollout.config import (
- RolloutClientSettings,
- RolloutServerSettings,
- RolloutSettings,
- configure,
- get_settings,
- reset_settings,
+from osmosis_ai.rollout.validator import (
+ AgentLoopValidationError,
+ ValidationError,
+ ValidationResult,
+ validate_agent_loop,
)
__all__ = [
- # Core classes
- "RolloutAgentLoop",
- "RolloutContext",
- "RolloutResult",
- # Client
- "OsmosisLLMClient",
- "CompletionsResult",
- "LLMClientProtocol",
- # Server
- "create_app",
- "serve_agent_loop",
- "validate_and_report",
- "ServeError",
"DEFAULT_HOST",
+ # Schemas - Configuration
+ "DEFAULT_MAX_METADATA_SIZE_BYTES",
"DEFAULT_PORT",
- # Validation
- "validate_agent_loop",
- "ValidationResult",
- "ValidationError",
- "AgentLoopValidationError",
+ "AgentLoopNotFoundError",
# Registry
"AgentLoopRegistry",
- "register_agent_loop",
- "unregister_agent_loop",
- "get_agent_loop",
- "list_agent_loops",
- # Schemas - Request/Response
- "RolloutRequest",
- "RolloutResponse",
- "InitResponse",
- "RolloutStatus",
- "RolloutMetrics",
+ "AgentLoopValidationError",
+ "CompletionUsage",
+ "CompletionsChoice",
# Schemas - Completions
"CompletionsRequest",
"CompletionsResponse",
- "CompletionsChoice",
- "CompletionUsage",
- # Schemas - Tool Definition
- "OpenAIFunctionToolSchema",
- "OpenAIFunctionSchema",
+ "CompletionsResult",
+ "InitResponse",
+ "LLMClientProtocol",
+ # Schemas - Type Aliases
+ "MessageDict",
+ "OpenAIFunctionCallSchema",
"OpenAIFunctionParametersSchema",
- "OpenAIFunctionPropertySchema",
# Schemas - Tool Call (adapted from verl)
"OpenAIFunctionParsedSchema",
- "OpenAIFunctionCallSchema",
+ "OpenAIFunctionPropertySchema",
+ "OpenAIFunctionSchema",
"OpenAIFunctionToolCall",
- "ToolResponse",
- # Schemas - Type Aliases
- "MessageDict",
- "SamplingParamsDict",
- # Schemas - Configuration
- "DEFAULT_MAX_METADATA_SIZE_BYTES",
- "get_max_metadata_size_bytes",
- "set_max_metadata_size_bytes",
+ # Schemas - Tool Definition
+ "OpenAIFunctionToolSchema",
+ # Client
+ "OsmosisLLMClient",
# Exceptions
"OsmosisRolloutError",
- "OsmosisTransportError",
"OsmosisServerError",
- "OsmosisValidationError",
"OsmosisTimeoutError",
- "AgentLoopNotFoundError",
- "ToolExecutionError",
+ "OsmosisTransportError",
+ "OsmosisValidationError",
+ "PublicIPDetectionError",
+ # Core classes
+ "RolloutAgentLoop",
+ "RolloutClientSettings",
+ "RolloutContext",
+ "RolloutMetrics",
+ # Schemas - Request/Response
+ "RolloutRequest",
+ "RolloutResponse",
+ "RolloutResult",
+ "RolloutServerSettings",
+ # Configuration
+ "RolloutSettings",
+ "RolloutStatus",
+ "SamplingParamsDict",
+ "ServeError",
"ToolArgumentError",
+ "ToolExecutionError",
+ "ToolResponse",
+ "ValidationError",
+ "ValidationResult",
+ "configure",
+ "count_messages_by_role",
+ # Server
+ "create_app",
+ "create_tool_error_result",
# Tool utilities
"create_tool_result",
- "create_tool_error_result",
- "serialize_tool_result",
- "parse_tool_arguments",
- "get_tool_call_info",
+ # Network utilities
+ "detect_public_ip",
"execute_tool_calls",
- # Message utilities
- "parse_tool_calls",
- "normalize_stop",
+ "get_agent_loop",
+ "get_max_metadata_size_bytes",
"get_message_content",
"get_message_role",
+ "get_settings",
+ "get_tool_call_info",
"is_assistant_message",
+ "is_private_ip",
"is_tool_message",
"is_user_message",
- "count_messages_by_role",
- # Network utilities
- "detect_public_ip",
- "PublicIPDetectionError",
- "validate_ipv4",
"is_valid_hostname_or_ip",
- "is_private_ip",
- # Configuration
- "RolloutSettings",
- "RolloutClientSettings",
- "RolloutServerSettings",
- "get_settings",
- "configure",
+ "list_agent_loops",
+ "normalize_stop",
+ "parse_tool_arguments",
+ # Message utilities
+ "parse_tool_calls",
+ "register_agent_loop",
"reset_settings",
+ "serialize_tool_result",
+ "serve_agent_loop",
+ "set_max_metadata_size_bytes",
+ "unregister_agent_loop",
+ # Validation
+ "validate_agent_loop",
+ "validate_and_report",
+ "validate_ipv4",
]
diff --git a/osmosis_ai/rollout/_compat.py b/osmosis_ai/rollout/_compat.py
index 39684e1e..0c102d53 100644
--- a/osmosis_ai/rollout/_compat.py
+++ b/osmosis_ai/rollout/_compat.py
@@ -17,13 +17,13 @@
from __future__ import annotations
-from typing import Any, Optional, Tuple
+from typing import Any
def import_optional(
module_name: str,
- package_name: Optional[str] = None,
-) -> Tuple[Any, bool]:
+ package_name: str | None = None,
+) -> tuple[Any, bool]:
"""Attempt to import an optional module.
Args:
@@ -49,9 +49,9 @@ def import_optional(
def require_optional(
module_name: str,
- package_name: Optional[str] = None,
+ package_name: str | None = None,
feature_name: str = "",
- install_extra: Optional[str] = None,
+ install_extra: str | None = None,
) -> Any:
"""Import an optional module, raising a friendly error if unavailable.
@@ -100,15 +100,15 @@ def require_optional(
__all__ = [
- # Functions
- "import_optional",
- "require_optional",
+ "FASTAPI_AVAILABLE",
# Availability flags
"PYDANTIC_SETTINGS_AVAILABLE",
- "FASTAPI_AVAILABLE",
"UVICORN_AVAILABLE",
+ "fastapi",
+ # Functions
+ "import_optional",
# Pre-imported modules (may be None)
"pydantic_settings",
- "fastapi",
+ "require_optional",
"uvicorn",
]
diff --git a/osmosis_ai/rollout/cli.py b/osmosis_ai/rollout/cli.py
index 0c35ec13..6201c2d2 100644
--- a/osmosis_ai/rollout/cli.py
+++ b/osmosis_ai/rollout/cli.py
@@ -11,7 +11,15 @@
import argparse
import sys
-from typing import Optional
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from osmosis_ai.rollout.eval.evaluation.cli import (
+ EvalCommand as _EvalCommandImpl,
+ )
+ from osmosis_ai.rollout.eval.test_mode.cli import (
+ TestCommand as _TestCommandImpl,
+ )
from osmosis_ai.rollout.cli_utils import CLIError, load_agent_loop
from osmosis_ai.rollout.server.serve import (
@@ -20,8 +28,6 @@
serve_agent_loop,
validate_and_report,
)
-from osmosis_ai.rollout.validator import validate_agent_loop
-
# Re-export CLIError for backwards compatibility
# (load_agent_loop is private, renamed from _load_agent_loop)
@@ -228,12 +234,14 @@ class TestCommand:
def __init__(self) -> None:
"""Initialize with lazy-loaded implementation."""
- self._impl: Optional["_TestCommandImpl"] = None
+ self._impl: _TestCommandImpl | None = None
- def _get_impl(self) -> "_TestCommandImpl":
+ def _get_impl(self) -> _TestCommandImpl:
"""Lazily load the actual TestCommand implementation."""
if self._impl is None:
- from osmosis_ai.rollout.eval.test_mode.cli import TestCommand as _TestCommandImpl
+ from osmosis_ai.rollout.eval.test_mode.cli import (
+ TestCommand as _TestCommandImpl,
+ )
self._impl = _TestCommandImpl()
return self._impl
@@ -250,11 +258,13 @@ class EvalCommand:
"""
def __init__(self) -> None:
- self._impl: Optional["_EvalCommandImpl"] = None
+ self._impl: _EvalCommandImpl | None = None
- def _get_impl(self) -> "_EvalCommandImpl":
+ def _get_impl(self) -> _EvalCommandImpl:
if self._impl is None:
- from osmosis_ai.rollout.eval.evaluation.cli import EvalCommand as _EvalCommandImpl
+ from osmosis_ai.rollout.eval.evaluation.cli import (
+ EvalCommand as _EvalCommandImpl,
+ )
self._impl = _EvalCommandImpl()
return self._impl
diff --git a/osmosis_ai/rollout/cli_utils.py b/osmosis_ai/rollout/cli_utils.py
index 10c7ff41..8d9ff331 100644
--- a/osmosis_ai/rollout/cli_utils.py
+++ b/osmosis_ai/rollout/cli_utils.py
@@ -49,15 +49,15 @@ def load_agent_loop(module_path: str) -> RolloutAgentLoop:
try:
module = importlib.import_module(module_name)
except ImportError as e:
- raise CLIError(f"Cannot import module '{module_name}': {e}")
+ raise CLIError(f"Cannot import module '{module_name}': {e}") from e
try:
agent_loop = getattr(module, attr_name)
- except AttributeError:
+ except AttributeError as e:
raise CLIError(
f"Module '{module_name}' has no attribute '{attr_name}'. "
f"Available attributes: {[a for a in dir(module) if not a.startswith('_')]}"
- )
+ ) from e
# If it's a class, instantiate it
if isinstance(agent_loop, type):
@@ -68,7 +68,7 @@ def load_agent_loop(module_path: str) -> RolloutAgentLoop:
try:
agent_loop = agent_loop()
except Exception as e:
- raise CLIError(f"Cannot instantiate '{attr_name}': {e}")
+ raise CLIError(f"Cannot instantiate '{attr_name}': {e}") from e
# Validate it's a RolloutAgentLoop instance
if not isinstance(agent_loop, RolloutAgentLoop):
diff --git a/osmosis_ai/rollout/client.py b/osmosis_ai/rollout/client.py
index 87e7081f..9b6a1b58 100644
--- a/osmosis_ai/rollout/client.py
+++ b/osmosis_ai/rollout/client.py
@@ -27,7 +27,7 @@
import logging
import time
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any
import httpx
@@ -73,10 +73,10 @@ class CompletionsResult:
print(result.content)
"""
- message: Dict[str, Any]
- token_ids: List[int]
- logprobs: List[float]
- usage: Dict[str, int]
+ message: dict[str, Any]
+ token_ids: list[int]
+ logprobs: list[float]
+ usage: dict[str, int]
finish_reason: str
@property
@@ -85,12 +85,12 @@ def has_tool_calls(self) -> bool:
return bool(self.message.get("tool_calls"))
@property
- def tool_calls(self) -> List[Dict[str, Any]]:
+ def tool_calls(self) -> list[dict[str, Any]]:
"""Get tool calls from the response."""
return self.message.get("tool_calls", [])
@property
- def content(self) -> Optional[str]:
+ def content(self) -> str | None:
"""Get text content from the response."""
return self.message.get("content")
@@ -136,11 +136,11 @@ def __init__(
self,
server_url: str,
rollout_id: str,
- api_key: Optional[str] = None,
- timeout_seconds: Optional[float] = None,
- max_retries: Optional[int] = None,
- complete_rollout_retries: Optional[int] = None,
- settings: Optional[RolloutClientSettings] = None,
+ api_key: str | None = None,
+ timeout_seconds: float | None = None,
+ max_retries: int | None = None,
+ complete_rollout_retries: int | None = None,
+ settings: RolloutClientSettings | None = None,
):
"""Initialize the LLM client.
@@ -161,7 +161,9 @@ def __init__(
self.rollout_id = rollout_id
self.api_key = api_key
self.timeout_seconds = timeout_seconds or settings.timeout_seconds
- self.max_retries = max_retries if max_retries is not None else settings.max_retries
+ self.max_retries = (
+ max_retries if max_retries is not None else settings.max_retries
+ )
self.complete_rollout_retries = (
complete_rollout_retries
if complete_rollout_retries is not None
@@ -175,7 +177,7 @@ def __init__(
self._retry_max_delay = settings.retry_max_delay
# HTTP client (lazy initialized)
- self._client: Optional[httpx.AsyncClient] = None
+ self._client: httpx.AsyncClient | None = None
self._client_lock = asyncio.Lock()
# Metrics tracking
@@ -215,12 +217,12 @@ async def close(self) -> None:
def _calculate_retry_delay(self, attempt: int) -> float:
"""Calculate delay for retry with exponential backoff."""
- delay = self._retry_base_delay * (2 ** attempt)
+ delay = self._retry_base_delay * (2**attempt)
return min(delay, self._retry_max_delay)
async def chat_completions(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
temperature: float = 1.0,
top_p: float = 1.0,
max_tokens: int = 512,
@@ -265,7 +267,7 @@ async def chat_completions(
client = await self._get_client()
url = f"{self.server_url}/v1/chat/completions"
- last_error: Optional[Exception] = None
+ last_error: Exception | None = None
for attempt in range(self.max_retries + 1):
start_time = time.monotonic()
@@ -285,7 +287,9 @@ async def chat_completions(
self._num_llm_calls += 1
if "usage" in data:
self._prompt_tokens += data["usage"].get("prompt_tokens", 0)
- self._response_tokens += data["usage"].get("completion_tokens", 0)
+ self._response_tokens += data["usage"].get(
+ "completion_tokens", 0
+ )
# Extract response
choice = data["choices"][0]
@@ -331,7 +335,9 @@ async def chat_completions(
last_error = OsmosisTimeoutError(str(e))
except httpx.RequestError as e:
- error_detail = f"{type(e).__name__}: {e}" if str(e) else f"{type(e).__name__}"
+ error_detail = (
+ f"{type(e).__name__}: {e}" if str(e) else f"{type(e).__name__}"
+ )
logger.warning(
"Transport error: attempt=%d/%d, error=%s",
attempt + 1,
@@ -360,11 +366,11 @@ async def chat_completions(
async def complete_rollout(
self,
status: str,
- final_messages: List[Dict[str, Any]],
+ final_messages: list[dict[str, Any]],
finish_reason: str = "stop",
- error_message: Optional[str] = None,
- metrics: Optional[RolloutMetrics] = None,
- reward: Optional[float] = None,
+ error_message: str | None = None,
+ metrics: RolloutMetrics | None = None,
+ reward: float | None = None,
) -> None:
"""Notify TrainGate that rollout is complete.
@@ -405,7 +411,7 @@ async def complete_rollout(
client = await self._get_client()
url = f"{self.server_url}/v1/rollout/completed"
- last_error: Optional[Exception] = None
+ last_error: Exception | None = None
for attempt in range(self.complete_rollout_retries + 1):
try:
@@ -451,7 +457,9 @@ async def complete_rollout(
last_error = OsmosisTimeoutError(str(e))
except httpx.RequestError as e:
- error_detail = f"{type(e).__name__}: {e}" if str(e) else f"{type(e).__name__}"
+ error_detail = (
+ f"{type(e).__name__}: {e}" if str(e) else f"{type(e).__name__}"
+ )
logger.warning(
"Completion transport error: attempt=%d, error=%s",
attempt + 1,
@@ -490,7 +498,7 @@ def get_metrics(self) -> RolloutMetrics:
response_tokens=self._response_tokens,
)
- async def __aenter__(self) -> "OsmosisLLMClient":
+ async def __aenter__(self) -> OsmosisLLMClient:
"""Async context manager entry."""
return self
diff --git a/osmosis_ai/rollout/config/__init__.py b/osmosis_ai/rollout/config/__init__.py
index 1b621257..7d676857 100644
--- a/osmosis_ai/rollout/config/__init__.py
+++ b/osmosis_ai/rollout/config/__init__.py
@@ -27,12 +27,12 @@
)
__all__ = [
- # Settings classes
- "RolloutSettings",
"RolloutClientSettings",
"RolloutServerSettings",
+ # Settings classes
+ "RolloutSettings",
+ "configure",
# Functions
"get_settings",
- "configure",
"reset_settings",
]
diff --git a/osmosis_ai/rollout/config/settings.py b/osmosis_ai/rollout/config/settings.py
index 2917f949..40525372 100644
--- a/osmosis_ai/rollout/config/settings.py
+++ b/osmosis_ai/rollout/config/settings.py
@@ -19,16 +19,12 @@
from __future__ import annotations
-from typing import Optional
-
from pydantic import BaseModel, Field
from osmosis_ai.rollout._compat import (
PYDANTIC_SETTINGS_AVAILABLE,
- pydantic_settings,
)
-
# Conditionally use pydantic-settings BaseSettings or fallback to BaseModel
if PYDANTIC_SETTINGS_AVAILABLE:
from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -226,7 +222,7 @@ class RolloutSettings(_BaseSettings):
# Global settings singleton
-_settings: Optional[RolloutSettings] = None
+_settings: RolloutSettings | None = None
def get_settings() -> RolloutSettings:
@@ -280,8 +276,8 @@ def reset_settings() -> None:
"RolloutClientSettings",
"RolloutServerSettings",
"RolloutSettings",
+ "configure",
# Functions
"get_settings",
- "configure",
"reset_settings",
]
diff --git a/osmosis_ai/rollout/console.py b/osmosis_ai/rollout/console.py
index cf7aea2d..d49331d7 100644
--- a/osmosis_ai/rollout/console.py
+++ b/osmosis_ai/rollout/console.py
@@ -16,17 +16,18 @@
from __future__ import annotations
import sys
-from typing import Any, Callable, List, Optional
+from collections.abc import Callable
+from typing import Any
# Try to import rich, gracefully degrade if not available
try:
+ from rich import box
from rich.console import Console as RichConsole
from rich.markup import escape as rich_escape
from rich.panel import Panel
+ from rich.style import Style
from rich.table import Table
from rich.text import Text
- from rich.style import Style
- from rich import box
RICH_AVAILABLE = True
except ImportError:
@@ -83,7 +84,7 @@ def __init__(
self,
*,
file: Any = None,
- force_terminal: Optional[bool] = None,
+ force_terminal: bool | None = None,
no_color: bool = False,
):
"""Initialize the console.
@@ -140,7 +141,7 @@ def run_rich(self, render_fn: Callable[[Any], None]) -> bool:
except ImportError:
return False
- def _get_ansi_style(self, style: Optional[str]) -> str:
+ def _get_ansi_style(self, style: str | None) -> str:
"""Get ANSI escape code for a style name."""
if not style or self._no_color or not self._is_tty:
return ""
@@ -154,7 +155,7 @@ def _filter_plain_print_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
def print(
self,
*args: Any,
- style: Optional[str] = None,
+ style: str | None = None,
end: str = "\n",
**kwargs: Any,
) -> None:
@@ -265,7 +266,10 @@ def panel(
# Content lines
for line in lines:
padded = line.ljust(box_width - 4)
- print(f"{style_code}│{reset} {padded} {style_code}│{reset}", file=self._file)
+ print(
+ f"{style_code}│{reset} {padded} {style_code}│{reset}",
+ file=self._file,
+ )
# Bottom border
bottom = f"╰{'─' * (box_width - 2)}╯"
@@ -273,10 +277,10 @@ def panel(
def table(
self,
- rows: List[tuple[str, str]],
+ rows: list[tuple[str, str]],
*,
- title: Optional[str] = None,
- headers: Optional[tuple[str, str]] = None,
+ title: str | None = None,
+ headers: tuple[str, str] | None = None,
) -> None:
"""Print a simple two-column table.
@@ -346,7 +350,7 @@ def format_styled(self, text: str, style: str) -> str:
return f"{ansi_style}{text}{_AnsiColors.RESET}"
return text
- def input(self, prompt: str = "", style: Optional[str] = None) -> str:
+ def input(self, prompt: str = "", style: str | None = None) -> str:
"""Get user input with optional styled prompt.
Args:
@@ -372,7 +376,7 @@ def input(self, prompt: str = "", style: Optional[str] = None) -> str:
__all__ = [
+ "RICH_AVAILABLE",
"Console",
"console",
- "RICH_AVAILABLE",
]
diff --git a/osmosis_ai/rollout/core/__init__.py b/osmosis_ai/rollout/core/__init__.py
index 721f678a..bf8544d9 100644
--- a/osmosis_ai/rollout/core/__init__.py
+++ b/osmosis_ai/rollout/core/__init__.py
@@ -21,7 +21,6 @@
RolloutContext,
RolloutResult,
)
-from osmosis_ai.rollout.core.llm_client import LLMClientProtocol
from osmosis_ai.rollout.core.exceptions import (
AgentLoopNotFoundError,
OsmosisRolloutError,
@@ -32,12 +31,13 @@
ToolArgumentError,
ToolExecutionError,
)
+from osmosis_ai.rollout.core.llm_client import LLMClientProtocol
from osmosis_ai.rollout.core.schemas import (
- CompletionUsage,
+ DEFAULT_MAX_METADATA_SIZE_BYTES,
CompletionsChoice,
CompletionsRequest,
CompletionsResponse,
- DEFAULT_MAX_METADATA_SIZE_BYTES,
+ CompletionUsage,
InitResponse,
MessageDict,
OpenAIFunctionCallSchema,
@@ -58,47 +58,47 @@
)
__all__ = [
- # Base classes
- "RolloutAgentLoop",
- "RolloutContext",
- "RolloutResult",
+ # Configuration
+ "DEFAULT_MAX_METADATA_SIZE_BYTES",
+ "AgentLoopNotFoundError",
+ "CompletionUsage",
+ "CompletionsChoice",
+ # Schemas - Completions
+ "CompletionsRequest",
+ "CompletionsResponse",
+ "InitResponse",
# Protocol
"LLMClientProtocol",
+ # Type aliases
+ "MessageDict",
+ "OpenAIFunctionCallSchema",
+ "OpenAIFunctionParametersSchema",
+ # Schemas - Tool Call (adapted from verl)
+ "OpenAIFunctionParsedSchema",
+ "OpenAIFunctionPropertySchema",
+ "OpenAIFunctionSchema",
+ "OpenAIFunctionToolCall",
+ # Schemas - Tool Definition
+ "OpenAIFunctionToolSchema",
# Exceptions
"OsmosisRolloutError",
- "OsmosisTransportError",
"OsmosisServerError",
- "OsmosisValidationError",
"OsmosisTimeoutError",
- "AgentLoopNotFoundError",
- "ToolExecutionError",
- "ToolArgumentError",
+ "OsmosisTransportError",
+ "OsmosisValidationError",
+ # Base classes
+ "RolloutAgentLoop",
+ "RolloutContext",
+ "RolloutMetrics",
# Schemas - Request/Response
"RolloutRequest",
"RolloutResponse",
- "InitResponse",
+ "RolloutResult",
"RolloutStatus",
- "RolloutMetrics",
- # Schemas - Completions
- "CompletionsRequest",
- "CompletionsResponse",
- "CompletionsChoice",
- "CompletionUsage",
- # Schemas - Tool Definition
- "OpenAIFunctionToolSchema",
- "OpenAIFunctionSchema",
- "OpenAIFunctionParametersSchema",
- "OpenAIFunctionPropertySchema",
- # Schemas - Tool Call (adapted from verl)
- "OpenAIFunctionParsedSchema",
- "OpenAIFunctionCallSchema",
- "OpenAIFunctionToolCall",
- "ToolResponse",
- # Type aliases
- "MessageDict",
"SamplingParamsDict",
- # Configuration
- "DEFAULT_MAX_METADATA_SIZE_BYTES",
+ "ToolArgumentError",
+ "ToolExecutionError",
+ "ToolResponse",
"get_max_metadata_size_bytes",
"set_max_metadata_size_bytes",
]
diff --git a/osmosis_ai/rollout/core/base.py b/osmosis_ai/rollout/core/base.py
index 74fda296..789d8650 100644
--- a/osmosis_ai/rollout/core/base.py
+++ b/osmosis_ai/rollout/core/base.py
@@ -35,7 +35,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.core.schemas import (
OpenAIFunctionToolSchema,
@@ -91,11 +91,11 @@ class RolloutResult:
"""
status: str
- final_messages: List[Dict[str, Any]]
+ final_messages: list[dict[str, Any]]
finish_reason: str
- error_message: Optional[str] = None
- metrics: Optional[RolloutMetrics] = None
- reward: Optional[float] = None
+ error_message: str | None = None
+ metrics: RolloutMetrics | None = None
+ reward: float | None = None
@dataclass
@@ -132,8 +132,8 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
"""
request: RolloutRequest
- tools: List[OpenAIFunctionToolSchema]
- llm: "LLMClientProtocol"
+ tools: list[OpenAIFunctionToolSchema]
+ llm: LLMClientProtocol
# Internal state for tracking metrics
_start_time: float = field(default=0.0, repr=False)
@@ -141,13 +141,13 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
_num_tool_calls: int = field(default=0, repr=False)
# Debug logging configuration
- _debug_dir: Optional[str] = field(default=None, repr=False)
+ _debug_dir: str | None = field(default=None, repr=False)
async def chat(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
**kwargs: Any,
- ) -> "CompletionsResult":
+ ) -> CompletionsResult:
"""Make a chat completion call to the LLM.
Shorthand for self.llm.chat_completions(). Passes through
@@ -170,9 +170,9 @@ async def chat(
def complete(
self,
- final_messages: List[Dict[str, Any]],
+ final_messages: list[dict[str, Any]],
finish_reason: str = "stop",
- reward: Optional[float] = None,
+ reward: float | None = None,
) -> RolloutResult:
"""Create a successful completion result.
@@ -203,7 +203,7 @@ def complete(
def error(
self,
error_message: str,
- final_messages: Optional[List[Dict[str, Any]]] = None,
+ final_messages: list[dict[str, Any]] | None = None,
) -> RolloutResult:
"""Create an error result.
@@ -249,7 +249,7 @@ def debug_enabled(self) -> bool:
"""Check if debug logging is enabled."""
return self._debug_dir is not None
- def _get_debug_file_path(self) -> Optional[str]:
+ def _get_debug_file_path(self) -> str | None:
"""Get the debug file path for this rollout."""
if not self._debug_dir:
return None
@@ -310,7 +310,9 @@ def log_event(self, event_type: str, **data: Any) -> None:
try:
os.makedirs(self._debug_dir, exist_ok=True)
except OSError as e:
- logger.warning("Failed to create debug directory %s: %s", self._debug_dir, e)
+ logger.warning(
+ "Failed to create debug directory %s: %s", self._debug_dir, e
+ )
return
# Build event entry
@@ -416,7 +418,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
name: str # Must be set by subclass
@abstractmethod
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
"""Return the tools available for this rollout.
Called when a new rollout request arrives. Can return different
diff --git a/osmosis_ai/rollout/core/exceptions.py b/osmosis_ai/rollout/core/exceptions.py
index 2e2123ec..a042950a 100644
--- a/osmosis_ai/rollout/core/exceptions.py
+++ b/osmosis_ai/rollout/core/exceptions.py
@@ -28,8 +28,6 @@
from __future__ import annotations
-from typing import List
-
class OsmosisRolloutError(Exception):
"""Base exception for all Osmosis rollout errors.
@@ -121,7 +119,7 @@ class AgentLoopNotFoundError(OsmosisRolloutError):
available: List of registered agent loop names.
"""
- def __init__(self, name: str, available: List[str]):
+ def __init__(self, name: str, available: list[str]):
"""Initialize with lookup name and available names.
Args:
@@ -130,9 +128,7 @@ def __init__(self, name: str, available: List[str]):
"""
self.name = name
self.available = available
- super().__init__(
- f"Agent loop '{name}' not found. Available: {available}"
- )
+ super().__init__(f"Agent loop '{name}' not found. Available: {available}")
class ToolExecutionError(OsmosisRolloutError):
diff --git a/osmosis_ai/rollout/core/llm_client.py b/osmosis_ai/rollout/core/llm_client.py
index 3c26df70..2c5aa5d3 100644
--- a/osmosis_ai/rollout/core/llm_client.py
+++ b/osmosis_ai/rollout/core/llm_client.py
@@ -6,7 +6,7 @@
from __future__ import annotations
-from typing import TYPE_CHECKING, Any, Dict, List, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
if TYPE_CHECKING:
from osmosis_ai.rollout.client import CompletionsResult
@@ -23,9 +23,9 @@ class LLMClientProtocol(Protocol):
async def chat_completions(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
**kwargs: Any,
- ) -> "CompletionsResult":
+ ) -> CompletionsResult:
"""Make a chat completion request.
Args:
@@ -37,7 +37,7 @@ async def chat_completions(
"""
...
- def get_metrics(self) -> "RolloutMetrics":
+ def get_metrics(self) -> RolloutMetrics:
"""Return accumulated metrics from this client session.
Returns:
diff --git a/osmosis_ai/rollout/core/schemas.py b/osmosis_ai/rollout/core/schemas.py
index 60b7cdab..27b9d1d3 100644
--- a/osmosis_ai/rollout/core/schemas.py
+++ b/osmosis_ai/rollout/core/schemas.py
@@ -31,17 +31,16 @@
import json
import threading
from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Literal
from urllib.parse import urlparse
from pydantic import BaseModel, Field, field_validator, model_validator
-
# =============================================================================
# Type Aliases
# =============================================================================
-MessageDict = Dict[str, Any]
+MessageDict = dict[str, Any]
"""Type alias for message dicts in protocol transmission.
Supports the full OpenAI message format including tool_call_id for tool responses.
@@ -50,7 +49,7 @@
{"role": "tool", "content": "345", "tool_call_id": "call_123"}
"""
-SamplingParamsDict = Dict[str, Any]
+SamplingParamsDict = dict[str, Any]
"""Type alias for sampling parameters dict.
Standard keys: temperature, top_p, max_tokens, stop, logprobs.
@@ -83,8 +82,8 @@ class OpenAIFunctionPropertySchema(BaseModel):
"""
type: str
- description: Optional[str] = None
- enum: Optional[List[str]] = None
+ description: str | None = None
+ enum: list[str] | None = None
class OpenAIFunctionParametersSchema(BaseModel):
@@ -97,8 +96,8 @@ class OpenAIFunctionParametersSchema(BaseModel):
"""
type: str
- properties: Dict[str, OpenAIFunctionPropertySchema]
- required: List[str]
+ properties: dict[str, OpenAIFunctionPropertySchema]
+ required: list[str]
class OpenAIFunctionSchema(BaseModel):
@@ -211,12 +210,12 @@ class OpenAIFunctionCallSchema(BaseModel):
"""
name: str
- arguments: Dict[str, Any]
+ arguments: dict[str, Any]
@staticmethod
def from_openai_function_parsed_schema(
parsed_schema: OpenAIFunctionParsedSchema,
- ) -> Tuple["OpenAIFunctionCallSchema", bool]:
+ ) -> tuple[OpenAIFunctionCallSchema, bool]:
"""Parse a function call from LLM output.
Args:
@@ -292,28 +291,34 @@ class ToolResponse(BaseModel):
)
"""
- text: Optional[str] = None
- image: Optional[List[Any]] = None
- video: Optional[List[Any]] = None
+ text: str | None = None
+ image: list[Any] | None = None
+ video: list[Any] | None = None
@model_validator(mode="before")
@classmethod
- def validate_media_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+ def validate_media_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
"""Validate that image and video fields are lists if provided."""
- if "image" in values and values["image"] is not None:
- if not isinstance(values["image"], list):
- raise ValueError(
- f"image must be a list, but got {type(values['image'])}. "
- f"For single images, wrap in a list: [image]. "
- f"Example: {{'image': [img1]}} or {{'image': [img1, img2, ...]}}."
- )
- if "video" in values and values["video"] is not None:
- if not isinstance(values["video"], list):
- raise ValueError(
- f"video must be a list, but got {type(values['video'])}. "
- f"For single videos, wrap in a list: [video]. "
- f"Example: {{'video': [video1]}} or {{'video': [video1, video2, ...]}}."
- )
+ if (
+ "image" in values
+ and values["image"] is not None
+ and not isinstance(values["image"], list)
+ ):
+ raise ValueError(
+ f"image must be a list, but got {type(values['image'])}. "
+ f"For single images, wrap in a list: [image]. "
+ f"Example: {{'image': [img1]}} or {{'image': [img1, img2, ...]}}."
+ )
+ if (
+ "video" in values
+ and values["video"] is not None
+ and not isinstance(values["video"], list)
+ ):
+ raise ValueError(
+ f"video must be a list, but got {type(values['video'])}. "
+ f"For single videos, wrap in a list: [video]. "
+ f"Example: {{'video': [video1]}} or {{'video': [video1, video2, ...]}}."
+ )
return values
def is_empty(self) -> bool:
@@ -454,14 +459,14 @@ class RolloutRequest(BaseModel):
rollout_id: str = Field(min_length=1, max_length=256)
server_url: str
- messages: List[MessageDict]
+ messages: list[MessageDict]
completion_params: SamplingParamsDict
- tool_server_url: Optional[str] = None
+ tool_server_url: str | None = None
max_turns: int = 10
max_tokens_total: int = 8192
- metadata: Dict[str, Any] = Field(default_factory=dict)
- api_key: Optional[str] = None
- idempotency_key: Optional[str] = None
+ metadata: dict[str, Any] = Field(default_factory=dict)
+ api_key: str | None = None
+ idempotency_key: str | None = None
@field_validator("rollout_id")
@classmethod
@@ -483,7 +488,7 @@ def validate_server_url_format(cls, v: str) -> str:
return v
@model_validator(mode="after")
- def validate_metadata_size(self) -> "RolloutRequest":
+ def validate_metadata_size(self) -> RolloutRequest:
"""Validate metadata size does not exceed limit."""
if self.metadata:
try:
@@ -497,7 +502,7 @@ def validate_metadata_size(self) -> "RolloutRequest":
except (TypeError, ValueError) as e:
if "exceeds maximum" in str(e):
raise
- raise ValueError(f"metadata must be JSON serializable: {e}")
+ raise ValueError(f"metadata must be JSON serializable: {e}") from e
return self
@@ -512,7 +517,7 @@ class InitResponse(BaseModel):
"""
rollout_id: str
- tools: List[OpenAIFunctionToolSchema] = Field(default_factory=list)
+ tools: list[OpenAIFunctionToolSchema] = Field(default_factory=list)
class RolloutResponse(BaseModel):
@@ -533,12 +538,12 @@ class RolloutResponse(BaseModel):
rollout_id: str
status: RolloutStatus
- final_messages: List[MessageDict] = Field(default_factory=list)
- finish_reason: Optional[str] = None
- error_message: Optional[str] = None
- reward: Optional[float] = None
- metrics: Optional[RolloutMetrics] = None
- extra_fields: Dict[str, Any] = Field(default_factory=dict)
+ final_messages: list[MessageDict] = Field(default_factory=list)
+ finish_reason: str | None = None
+ error_message: str | None = None
+ reward: float | None = None
+ metrics: RolloutMetrics | None = None
+ extra_fields: dict[str, Any] = Field(default_factory=dict)
# =============================================================================
@@ -566,12 +571,12 @@ class CompletionsRequest(BaseModel):
"""
model: str = "default"
- messages: List[MessageDict]
+ messages: list[MessageDict]
rollout_id: str = Field(min_length=1, max_length=256)
temperature: float = 1.0
top_p: float = 1.0
max_tokens: int = 512
- stop: Optional[List[str]] = None
+ stop: list[str] | None = None
logprobs: bool = True
@field_validator("rollout_id")
@@ -637,10 +642,10 @@ class CompletionsResponse(BaseModel):
object: str = "chat.completion"
created: int
model: str = "default"
- choices: List[CompletionsChoice]
- usage: Optional[CompletionUsage] = None
+ choices: list[CompletionsChoice]
+ usage: CompletionUsage | None = None
# TrainGate internal fields - NOT transmitted via HTTP to RolloutServer.
# These are accumulated in SessionManager for building AgentLoopOutput.
- token_ids: Optional[List[int]] = None
- logprobs: Optional[List[float]] = None
- prompt_token_ids: Optional[List[int]] = None
+ token_ids: list[int] | None = None
+ logprobs: list[float] | None = None
+ prompt_token_ids: list[int] | None = None
diff --git a/osmosis_ai/rollout/eval/common/__init__.py b/osmosis_ai/rollout/eval/common/__init__.py
index cdf364b7..ab95c1fa 100644
--- a/osmosis_ai/rollout/eval/common/__init__.py
+++ b/osmosis_ai/rollout/eval/common/__init__.py
@@ -31,25 +31,25 @@
)
__all__ = [
- "build_completion_params",
- "create_llm_client",
- "format_duration",
- "format_tokens",
- "load_agent",
- "load_dataset_rows",
"REQUIRED_COLUMNS",
+ "DatasetParseError",
"DatasetReader",
"DatasetRow",
- "dataset_row_to_request",
- "LocalExecutionError",
"DatasetValidationError",
- "DatasetParseError",
- "ToolValidationError",
- "ProviderError",
"ExternalLLMClient",
"LocalBatchResult",
- "LocalRunResult",
+ "LocalExecutionError",
"LocalLLMClientProtocol",
"LocalRolloutRunner",
+ "LocalRunResult",
+ "ProviderError",
+ "ToolValidationError",
+ "build_completion_params",
+ "create_llm_client",
+ "dataset_row_to_request",
+ "format_duration",
+ "format_tokens",
+ "load_agent",
+ "load_dataset_rows",
"validate_tools",
]
diff --git a/osmosis_ai/rollout/eval/common/cli.py b/osmosis_ai/rollout/eval/common/cli.py
index 4189129b..5d7cb550 100644
--- a/osmosis_ai/rollout/eval/common/cli.py
+++ b/osmosis_ai/rollout/eval/common/cli.py
@@ -3,7 +3,7 @@
from __future__ import annotations
import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.cli_utils import CLIError, load_agent_loop
from osmosis_ai.rollout.console import Console
@@ -16,7 +16,7 @@
)
# Mapping of LiteLLM provider prefixes to their expected environment variables.
-_PROVIDER_ENV_KEYS: Dict[str, str] = {
+_PROVIDER_ENV_KEYS: dict[str, str] = {
"openai": "OPENAI_API_KEY",
"anthropic": "ANTHROPIC_API_KEY",
"groq": "GROQ_API_KEY",
@@ -49,7 +49,7 @@ def format_duration(ms: float) -> str:
if ms < 1000:
return f"{ms:.0f}ms"
if ms < 60000:
- return f"{ms/1000:.1f}s"
+ return f"{ms / 1000:.1f}s"
minutes = int(ms // 60000)
seconds = (ms % 60000) / 1000
return f"{minutes}m{seconds:.1f}s"
@@ -64,7 +64,7 @@ def load_agent(
module: str,
quiet: bool,
console: Console,
-) -> Tuple[Optional["RolloutAgentLoop"], Optional[str]]:
+) -> tuple[RolloutAgentLoop | None, str | None]:
"""Load agent loop from module path."""
if not quiet:
console.print(f"Loading agent: {module}")
@@ -84,7 +84,7 @@ def load_mcp_agent(
mcp_path: str,
quiet: bool,
console: Console,
-) -> Tuple[Optional["RolloutAgentLoop"], Optional[str]]:
+) -> tuple[RolloutAgentLoop | None, str | None]:
"""Load an MCPAgentLoop from an MCP directory.
The directory must contain a ``main.py`` with a FastMCP instance and
@@ -94,8 +94,7 @@ def load_mcp_agent(
from osmosis_ai.rollout.mcp import MCPAgentLoop, MCPLoadError, load_mcp_server
except ImportError:
return None, (
- "MCP support requires fastmcp. "
- "Install it with: pip install osmosis-ai[mcp]"
+ "MCP support requires fastmcp. Install it with: pip install osmosis-ai[mcp]"
)
if not quiet:
@@ -110,20 +109,22 @@ def load_mcp_agent(
if not quiet:
tool_names = [t.function.name for t in agent_loop.get_tools(None)] # type: ignore[arg-type]
- console.print(f" Discovered {len(tool_names)} tool(s): {', '.join(tool_names)}")
+ console.print(
+ f" Discovered {len(tool_names)} tool(s): {', '.join(tool_names)}"
+ )
return agent_loop, None
def load_dataset_rows(
dataset_path: str,
- limit: Optional[int],
+ limit: int | None,
offset: int,
quiet: bool,
console: Console,
empty_error: str,
action_label: str,
-) -> Tuple[Optional[List["DatasetRow"]], Optional[str]]:
+) -> tuple[list[DatasetRow] | None, str | None]:
"""Load rows from dataset and provide command-specific messaging."""
if not quiet:
console.print(f"Loading dataset: {dataset_path}")
@@ -151,9 +152,9 @@ def load_dataset_rows(
def _check_api_key(
model: str,
- api_key: Optional[str],
- base_url: Optional[str],
-) -> Optional[str]:
+ api_key: str | None,
+ base_url: str | None,
+) -> str | None:
"""Return an error message if the required API key is missing, else None."""
if api_key or base_url:
return None
@@ -184,12 +185,11 @@ def _first_line(message: str) -> str:
return message.strip()
-def _format_model_error(model: str, base_url: Optional[str], detail: str) -> str:
+def _format_model_error(model: str, base_url: str | None, detail: str) -> str:
"""Format a concise model validation error with actionable guidance."""
if base_url:
return (
- "Invalid model for --base-url. "
- f"Received model='{model}'. Details: {detail}"
+ f"Invalid model for --base-url. Received model='{model}'. Details: {detail}"
)
return (
"Invalid LiteLLM model format. Use 'provider/model' "
@@ -198,7 +198,7 @@ def _format_model_error(model: str, base_url: Optional[str], detail: str) -> str
)
-def _check_model(model: str, base_url: Optional[str]) -> Optional[str]:
+def _check_model(model: str, base_url: str | None) -> str | None:
"""Return an error message if model format is invalid, else None."""
candidate = model.strip()
if not candidate:
@@ -212,7 +212,9 @@ def _check_model(model: str, base_url: Optional[str]) -> Optional[str]:
normalized_model = candidate if "/" in candidate else f"openai/{candidate}"
provider, model_name = normalized_model.split("/", 1)
if not provider.strip() or not model_name.strip():
- return _format_model_error(candidate, base_url, "missing provider or model name")
+ return _format_model_error(
+ candidate, base_url, "missing provider or model name"
+ )
try:
import litellm
@@ -252,11 +254,11 @@ def _check_model(model: str, base_url: Optional[str]) -> Optional[str]:
def create_llm_client(
model: str,
- api_key: Optional[str],
- base_url: Optional[str],
+ api_key: str | None,
+ base_url: str | None,
quiet: bool,
console: Console,
-) -> Tuple[Optional["ExternalLLMClient"], Optional[str]]:
+) -> tuple[ExternalLLMClient | None, str | None]:
"""Initialize ExternalLLMClient with consistent messaging and errors."""
if error := _check_model(model, base_url):
return None, error
@@ -290,10 +292,10 @@ def create_llm_client(
async def verify_llm_client(
- llm_client: "ExternalLLMClient",
+ llm_client: ExternalLLMClient,
quiet: bool,
console: Console,
-) -> Optional[str]:
+) -> str | None:
"""Run a preflight check against the LLM provider.
Returns an error message string on failure, or None on success.
@@ -310,17 +312,18 @@ async def verify_llm_client(
def build_completion_params(
- temperature: Optional[float],
- max_tokens: Optional[int],
-) -> Dict[str, Any]:
+ temperature: float | None,
+ max_tokens: int | None,
+) -> dict[str, Any]:
"""Build completion params dict from CLI options."""
- params: Dict[str, Any] = {}
+ params: dict[str, Any] = {}
if temperature is not None:
params["temperature"] = temperature
if max_tokens is not None:
params["max_tokens"] = max_tokens
return params
+
__all__ = [
"build_completion_params",
"create_llm_client",
diff --git a/osmosis_ai/rollout/eval/common/dataset.py b/osmosis_ai/rollout/eval/common/dataset.py
index de4bb1ac..58b6718c 100644
--- a/osmosis_ai/rollout/eval/common/dataset.py
+++ b/osmosis_ai/rollout/eval/common/dataset.py
@@ -9,11 +9,15 @@
import json
import logging
+from collections.abc import Iterator
from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, TypedDict, cast
+from typing import Any, TypedDict, cast
from osmosis_ai.rollout.core.schemas import RolloutRequest
-from osmosis_ai.rollout.eval.common.errors import DatasetParseError, DatasetValidationError
+from osmosis_ai.rollout.eval.common.errors import (
+ DatasetParseError,
+ DatasetValidationError,
+)
logger = logging.getLogger(__name__)
@@ -46,19 +50,19 @@ def __init__(self, file_path: str) -> None:
f"Supported formats: .json, .jsonl, .parquet"
)
- self._row_count: Optional[int] = None
+ self._row_count: int | None = None
def read(
self,
- limit: Optional[int] = None,
+ limit: int | None = None,
offset: int = 0,
- ) -> List[DatasetRow]:
+ ) -> list[DatasetRow]:
"""Read rows from the dataset file."""
return list(self.iter_rows(limit=limit, offset=offset))
def iter_rows(
self,
- limit: Optional[int] = None,
+ limit: int | None = None,
offset: int = 0,
) -> Iterator[DatasetRow]:
"""Yield validated rows lazily for memory-efficient iteration."""
@@ -78,7 +82,7 @@ def iter_rows(
yield validated
count += 1
- def _iter_raw_rows(self) -> Iterator[Dict[str, Any]]:
+ def _iter_raw_rows(self) -> Iterator[dict[str, Any]]:
if self._extension == ".json":
yield from self._parse_json()
elif self._extension == ".jsonl":
@@ -101,9 +105,9 @@ def __len__(self) -> int:
return self._row_count
- def _parse_json(self) -> List[Dict[str, Any]]:
+ def _parse_json(self) -> list[dict[str, Any]]:
try:
- with open(self.file_path, "r", encoding="utf-8") as f:
+ with open(self.file_path, encoding="utf-8") as f:
data = json.load(f)
except json.JSONDecodeError as e:
raise DatasetParseError(f"Invalid JSON file: {e}") from e
@@ -117,9 +121,9 @@ def _parse_json(self) -> List[Dict[str, Any]]:
return data
- def _iter_jsonl(self) -> Iterator[Dict[str, Any]]:
+ def _iter_jsonl(self) -> Iterator[dict[str, Any]]:
try:
- with open(self.file_path, "r", encoding="utf-8") as f:
+ with open(self.file_path, encoding="utf-8") as f:
for line_num, line in enumerate(f, start=1):
line = line.strip()
if not line:
@@ -137,7 +141,7 @@ def _iter_jsonl(self) -> Iterator[Dict[str, Any]]:
def _count_jsonl_rows(self) -> int:
count = 0
try:
- with open(self.file_path, "r", encoding="utf-8") as f:
+ with open(self.file_path, encoding="utf-8") as f:
for line in f:
if line.strip():
count += 1
@@ -145,7 +149,7 @@ def _count_jsonl_rows(self) -> int:
pass
return count
- def _parse_parquet(self) -> List[Dict[str, Any]]:
+ def _parse_parquet(self) -> list[dict[str, Any]]:
try:
import pyarrow.parquet as pq
except ImportError as e:
@@ -179,27 +183,33 @@ def _validate_row(self, row: Any, row_index: int) -> DatasetRow:
f"Row {row_index}: Expected object, got {type(row).__name__}"
)
- lower_keys = {k.lower(): k for k in row.keys()}
+ lower_keys = {k.lower(): k for k in row}
- missing = [required for required in REQUIRED_COLUMNS if required not in lower_keys]
+ missing = [
+ required for required in REQUIRED_COLUMNS if required not in lower_keys
+ ]
if missing:
raise DatasetValidationError(
f"Row {row_index}: Missing required columns: {missing}"
)
- result: Dict[str, Any] = {}
+ result: dict[str, Any] = {}
for required in REQUIRED_COLUMNS:
original_key = lower_keys[required]
value = row[original_key]
if value is None:
- raise DatasetValidationError(f"Row {row_index}: '{required}' cannot be null")
+ raise DatasetValidationError(
+ f"Row {row_index}: '{required}' cannot be null"
+ )
if not isinstance(value, str):
raise DatasetValidationError(
f"Row {row_index}: '{required}' must be a string, "
f"got {type(value).__name__}"
)
if not value.strip():
- raise DatasetValidationError(f"Row {row_index}: '{required}' cannot be empty")
+ raise DatasetValidationError(
+ f"Row {row_index}: '{required}' cannot be empty"
+ )
result[required] = value
for key, value in row.items():
@@ -214,13 +224,13 @@ def dataset_row_to_request(
row_index: int,
max_turns: int = 10,
max_tokens_total: int = 4096,
- completion_params: Optional[Dict[str, Any]] = None,
+ completion_params: dict[str, Any] | None = None,
rollout_id_prefix: str = "local",
- rollout_id: Optional[str] = None,
- metadata_overrides: Optional[Dict[str, Any]] = None,
+ rollout_id: str | None = None,
+ metadata_overrides: dict[str, Any] | None = None,
) -> RolloutRequest:
"""Convert a dataset row to a local RolloutRequest."""
- metadata: Dict[str, Any] = {
+ metadata: dict[str, Any] = {
"ground_truth": row["ground_truth"],
"row_index": row_index,
}
@@ -247,6 +257,7 @@ def dataset_row_to_request(
metadata=metadata,
)
+
__all__ = [
"REQUIRED_COLUMNS",
"DatasetReader",
diff --git a/osmosis_ai/rollout/eval/common/errors.py b/osmosis_ai/rollout/eval/common/errors.py
index 474cc109..9aa45997 100644
--- a/osmosis_ai/rollout/eval/common/errors.py
+++ b/osmosis_ai/rollout/eval/common/errors.py
@@ -44,11 +44,12 @@ class SystemicProviderError(ProviderError):
pass
+
__all__ = [
- "LocalExecutionError",
- "DatasetValidationError",
"DatasetParseError",
- "ToolValidationError",
+ "DatasetValidationError",
+ "LocalExecutionError",
"ProviderError",
"SystemicProviderError",
+ "ToolValidationError",
]
diff --git a/osmosis_ai/rollout/eval/common/llm_client.py b/osmosis_ai/rollout/eval/common/llm_client.py
index 8a32bb4d..448335a8 100644
--- a/osmosis_ai/rollout/eval/common/llm_client.py
+++ b/osmosis_ai/rollout/eval/common/llm_client.py
@@ -5,11 +5,11 @@
from __future__ import annotations
-import logging
import inspect
+import logging
import time
import warnings
-from typing import Any, Dict, List, Optional
+from typing import Any
warnings.filterwarnings(
"ignore",
@@ -39,7 +39,7 @@ def _get_provider_message(e: Exception) -> str:
msg = getattr(e, "message", str(e))
prefix = f"litellm.{type(e).__name__}: "
if msg.startswith(prefix):
- msg = msg[len(prefix):]
+ msg = msg[len(prefix) :]
# Strip litellm noise after common markers
for marker in _NOISE_MARKERS:
@@ -90,7 +90,7 @@ def _first_line(message: str) -> str:
def _format_provider_hint(
model: str,
- api_base: Optional[str],
+ api_base: str | None,
raw_message: str,
) -> str:
"""Build a concise actionable error for provider/model format issues."""
@@ -136,8 +136,8 @@ class ExternalLLMClient:
def __init__(
self,
model: str = "gpt-5-mini",
- api_key: Optional[str] = None,
- api_base: Optional[str] = None,
+ api_key: str | None = None,
+ api_base: str | None = None,
) -> None:
self._litellm = _get_litellm()
@@ -164,7 +164,7 @@ def __init__(
self._api_key = api_key
self._api_base = api_base
- self._tools: Optional[List[Dict[str, Any]]] = None
+ self._tools: list[dict[str, Any]] | None = None
self._llm_latency_ms: float = 0.0
self._num_llm_calls: int = 0
@@ -172,7 +172,7 @@ def __init__(
self._response_tokens: int = 0
self._closed = False
- def set_tools(self, tools: List[Any]) -> None:
+ def set_tools(self, tools: list[Any]) -> None:
"""Set tools for the current execution row."""
if tools:
self._tools = [
@@ -195,13 +195,13 @@ def reset_metrics(self) -> None:
async def chat_completions(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
**kwargs: Any,
) -> CompletionsResult:
"""Make a chat completion request via LiteLLM."""
start_time = time.monotonic()
- request_kwargs: Dict[str, Any] = {
+ request_kwargs: dict[str, Any] = {
"model": self.model,
"messages": messages,
}
@@ -295,9 +295,7 @@ async def preflight_check(self) -> None:
self._BudgetExceededError,
self._APIConnectionError,
) as e:
- raise SystemicProviderError(
- _get_provider_message(e)
- ) from e
+ raise SystemicProviderError(_get_provider_message(e)) from e
except Exception as e:
classified = self._classify_unknown_error(e)
if isinstance(classified, SystemicProviderError):
@@ -319,9 +317,7 @@ def _classify_unknown_error(self, e: Exception) -> ProviderError:
# Systemic: auth / forbidden / model-not-found
if status_code in (401, 403, 404):
- return SystemicProviderError(
- f"LLM API error (HTTP {status_code}): {msg}"
- )
+ return SystemicProviderError(f"LLM API error (HTTP {status_code}): {msg}")
# Systemic: connection errors wrapped as InternalServerError etc.
if _is_connection_error_message(msg):
@@ -373,10 +369,11 @@ async def close(self) -> None:
except Exception as exc:
logger.debug("LiteLLM async client cleanup failed: %s", exc)
- async def __aenter__(self) -> "ExternalLLMClient":
+ async def __aenter__(self) -> ExternalLLMClient:
return self
async def __aexit__(self, *args: Any) -> None:
await self.close()
+
__all__ = ["ExternalLLMClient"]
diff --git a/osmosis_ai/rollout/eval/common/runner.py b/osmosis_ai/rollout/eval/common/runner.py
index e37ebb61..346f91a3 100644
--- a/osmosis_ai/rollout/eval/common/runner.py
+++ b/osmosis_ai/rollout/eval/common/runner.py
@@ -10,13 +10,17 @@
import logging
import re
import time
+from collections.abc import Callable
from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, Protocol
+from typing import Any, Protocol
from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext, RolloutResult
from osmosis_ai.rollout.core.schemas import OpenAIFunctionToolSchema
from osmosis_ai.rollout.eval.common.dataset import DatasetRow, dataset_row_to_request
-from osmosis_ai.rollout.eval.common.errors import SystemicProviderError, ToolValidationError
+from osmosis_ai.rollout.eval.common.errors import (
+ SystemicProviderError,
+ ToolValidationError,
+)
logger = logging.getLogger(__name__)
@@ -24,7 +28,7 @@
class LocalLLMClientProtocol(Protocol):
"""LLM client contract required by LocalRolloutRunner."""
- def set_tools(self, tools: List[Any]) -> None: ...
+ def set_tools(self, tools: list[Any]) -> None: ...
def clear_tools(self) -> None: ...
@@ -39,30 +43,30 @@ class LocalRunResult:
row_index: int
success: bool
- result: Optional[RolloutResult] = None
- error: Optional[str] = None
+ result: RolloutResult | None = None
+ error: str | None = None
duration_ms: float = 0.0
- token_usage: Dict[str, Any] = field(default_factory=dict)
+ token_usage: dict[str, Any] = field(default_factory=dict)
@dataclass
class LocalBatchResult:
"""Aggregated results from running a batch of rows."""
- results: List[LocalRunResult]
+ results: list[LocalRunResult]
total: int
passed: int
failed: int
total_duration_ms: float
total_tokens: int
stopped_early: bool = False
- stop_reason: Optional[str] = None
+ stop_reason: str | None = None
TOOL_NAME_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
-def validate_tools(tools: List[OpenAIFunctionToolSchema]) -> None:
+def validate_tools(tools: list[OpenAIFunctionToolSchema]) -> None:
"""Validate tool schemas before sending to provider APIs."""
for i, tool in enumerate(tools):
if not tool.function:
@@ -94,9 +98,9 @@ def __init__(
agent_loop: RolloutAgentLoop,
llm_client: LocalLLMClientProtocol,
debug: bool = False,
- debug_dir: Optional[str] = None,
+ debug_dir: str | None = None,
rollout_id_prefix: str = "local",
- request_metadata: Optional[Dict[str, Any]] = None,
+ request_metadata: dict[str, Any] | None = None,
) -> None:
self.agent_loop = agent_loop
self.llm_client = llm_client
@@ -105,7 +109,7 @@ def __init__(
self.request_metadata = dict(request_metadata or {})
if debug_dir is not None:
- self.debug_dir: Optional[str] = debug_dir
+ self.debug_dir: str | None = debug_dir
elif debug:
self.debug_dir = "./local_debug"
else:
@@ -116,9 +120,9 @@ async def run_single(
row: DatasetRow,
row_index: int,
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
- rollout_id: Optional[str] = None,
- request_metadata: Optional[Dict[str, Any]] = None,
+ completion_params: dict[str, Any] | None = None,
+ rollout_id: str | None = None,
+ request_metadata: dict[str, Any] | None = None,
) -> LocalRunResult:
"""Run a single dataset row."""
overall_start = time.monotonic()
@@ -185,14 +189,13 @@ async def run_single(
# accurate stats for the failing run.
try:
metrics = self.llm_client.get_metrics()
- tokens = (
- int(getattr(metrics, "prompt_tokens", 0))
- + int(getattr(metrics, "response_tokens", 0))
+ tokens = int(getattr(metrics, "prompt_tokens", 0)) + int(
+ getattr(metrics, "response_tokens", 0)
)
except Exception:
tokens = 0
- setattr(e, "duration_ms", (time.monotonic() - overall_start) * 1000)
- setattr(e, "tokens", tokens)
+ e.duration_ms = (time.monotonic() - overall_start) * 1000
+ e.tokens = tokens
raise
except Exception as e:
@@ -212,17 +215,17 @@ async def run_single(
async def run_batch(
self,
- rows: List[DatasetRow],
+ rows: list[DatasetRow],
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
- on_progress: Optional[Callable[[int, int, LocalRunResult], None]] = None,
+ completion_params: dict[str, Any] | None = None,
+ on_progress: Callable[[int, int, LocalRunResult], None] | None = None,
start_index: int = 0,
) -> LocalBatchResult:
"""Run multiple rows sequentially."""
- results: List[LocalRunResult] = []
+ results: list[LocalRunResult] = []
total_start = time.monotonic()
stopped_early = False
- stop_reason: Optional[str] = None
+ stop_reason: str | None = None
for i, row in enumerate(rows):
row_index = start_index + i
@@ -273,6 +276,7 @@ async def run_batch(
stop_reason=stop_reason,
)
+
__all__ = [
"LocalBatchResult",
"LocalLLMClientProtocol",
diff --git a/osmosis_ai/rollout/eval/evaluation/__init__.py b/osmosis_ai/rollout/eval/evaluation/__init__.py
index f8a16f98..c53ce847 100644
--- a/osmosis_ai/rollout/eval/evaluation/__init__.py
+++ b/osmosis_ai/rollout/eval/evaluation/__init__.py
@@ -12,21 +12,21 @@
EvalModelSummary,
EvalResult,
EvalRowResult,
- EvalRunResult,
EvalRunner,
+ EvalRunResult,
)
__all__ = [
+ "EvalEvalSummary",
"EvalFnError",
"EvalFnWrapper",
- "load_eval_fn",
- "load_eval_fns",
- "EvalRunResult",
- "EvalRowResult",
- "EvalEvalSummary",
"EvalModelSummary",
"EvalResult",
+ "EvalRowResult",
+ "EvalRunResult",
"EvalRunner",
- "pass_at_k",
"format_eval_report",
+ "load_eval_fn",
+ "load_eval_fns",
+ "pass_at_k",
]
diff --git a/osmosis_ai/rollout/eval/evaluation/cli.py b/osmosis_ai/rollout/eval/evaluation/cli.py
index afb4084f..1f1078e9 100644
--- a/osmosis_ai/rollout/eval/evaluation/cli.py
+++ b/osmosis_ai/rollout/eval/evaluation/cli.py
@@ -10,7 +10,7 @@
import json
import logging
from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.console import Console
from osmosis_ai.rollout.eval.common.cli import (
@@ -222,7 +222,7 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None:
def run(self, args: argparse.Namespace) -> int:
return asyncio.run(self._run_async(args))
- def _validate_args(self, args: argparse.Namespace) -> Optional[str]:
+ def _validate_args(self, args: argparse.Namespace) -> str | None:
if args.module and args.mcp:
return "--module and --mcp are mutually exclusive."
if not args.module and not args.mcp:
@@ -262,8 +262,11 @@ def _print_header(self, args: argparse.Namespace) -> None:
def _load_eval_fns(
self, args: argparse.Namespace
- ) -> Tuple[Optional[List["EvalFnWrapper"]], Optional[str]]:
- from osmosis_ai.rollout.eval.evaluation.eval_fn import EvalFnError, load_eval_fns
+ ) -> tuple[list[EvalFnWrapper] | None, str | None]:
+ from osmosis_ai.rollout.eval.evaluation.eval_fn import (
+ EvalFnError,
+ load_eval_fns,
+ )
if not args.quiet:
self.console.print(f"Loading eval functions: {', '.join(args.eval_fns)}")
@@ -279,11 +282,11 @@ def _load_eval_fns(
return eval_fns, None
- def _write_output(self, args: argparse.Namespace, result: "EvalResult") -> None:
+ def _write_output(self, args: argparse.Namespace, result: EvalResult) -> None:
if not args.output:
return
- config: Dict[str, Any] = {
+ config: dict[str, Any] = {
"model": args.model,
"n_runs": args.n_runs,
"pass_threshold": args.pass_threshold,
@@ -292,7 +295,7 @@ def _write_output(self, args: argparse.Namespace, result: "EvalResult") -> None:
if args.baseline_model:
config["baseline_model"] = args.baseline_model
- summary: Dict[str, Any] = {
+ summary: dict[str, Any] = {
"total_rows": result.total_rows,
"total_runs": result.total_runs,
"stopped_early": result.stopped_early,
@@ -311,7 +314,7 @@ def _write_output(self, args: argparse.Namespace, result: "EvalResult") -> None:
"total_duration_ms": result.total_duration_ms,
}
- output_data: Dict[str, Any] = {
+ output_data: dict[str, Any] = {
"config": config,
"summary": summary,
"rows": [
@@ -454,7 +457,9 @@ async def _run_async(self, args: argparse.Namespace) -> int:
return 1
assert baseline_llm_client is not None
- error = await verify_llm_client(baseline_llm_client, args.quiet, self.console)
+ error = await verify_llm_client(
+ baseline_llm_client, args.quiet, self.console
+ )
if error:
self.console.print_error(f"Error (baseline): {error}")
await baseline_llm_client.close()
@@ -471,7 +476,7 @@ async def _run_async(self, args: argparse.Namespace) -> int:
baseline_llm_client=baseline_llm_client,
)
- def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
+ def on_progress(current: int, total: int, result: EvalRunResult) -> None:
if args.quiet:
return
@@ -491,7 +496,9 @@ def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
error_suffix = ""
if not result.success and result.error:
error_text = result.error.replace("\n", " ")
- error_msg = error_text[:47] + "..." if len(error_text) > 50 else error_text
+ error_msg = (
+ error_text[:47] + "..." if len(error_text) > 50 else error_text
+ )
error_suffix = f" - {error_msg}"
status_styled = self.console.format_styled(status, status_style)
@@ -503,7 +510,9 @@ def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
if not args.quiet:
self.console.print()
n_info = f" x{args.n_runs} runs" if args.n_runs > 1 else ""
- batch_info = f", batch_size={args.batch_size}" if args.batch_size > 1 else ""
+ batch_info = (
+ f", batch_size={args.batch_size}" if args.batch_size > 1 else ""
+ )
model_info = " x2 models" if args.baseline_model else ""
self.console.print(
f"Running evaluation ({len(rows)} rows{n_info}{model_info}{batch_info})..."
@@ -517,7 +526,9 @@ def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
rows=rows,
n_runs=args.n_runs,
max_turns=args.max_turns,
- completion_params=completion_params if completion_params else None,
+ completion_params=completion_params
+ if completion_params
+ else None,
pass_threshold=args.pass_threshold,
on_progress=on_progress,
start_index=args.offset,
@@ -528,7 +539,9 @@ def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
rows=rows,
n_runs=args.n_runs,
max_turns=args.max_turns,
- completion_params=completion_params if completion_params else None,
+ completion_params=completion_params
+ if completion_params
+ else None,
pass_threshold=args.pass_threshold,
on_progress=on_progress,
start_index=args.offset,
@@ -550,4 +563,5 @@ def on_progress(current: int, total: int, result: "EvalRunResult") -> None:
self._write_output(args, eval_result)
return 0
+
__all__ = ["EvalCommand"]
diff --git a/osmosis_ai/rollout/eval/evaluation/eval_fn.py b/osmosis_ai/rollout/eval/evaluation/eval_fn.py
index 022fbeaa..c3305198 100644
--- a/osmosis_ai/rollout/eval/evaluation/eval_fn.py
+++ b/osmosis_ai/rollout/eval/evaluation/eval_fn.py
@@ -17,7 +17,8 @@ def my_eval(messages: list, ground_truth: str, metadata: dict, **kwargs) -> floa
import inspect
import os
import sys
-from typing import Any, Callable, Dict, List
+from collections.abc import Callable
+from typing import Any
class EvalFnError(Exception):
@@ -48,8 +49,8 @@ def _detect_mode(fn: Callable) -> str:
params = list(sig.parameters.keys())
if not params:
raise EvalFnError(
- f"Eval function has no parameters. "
- f"Expected (solution_str, ...) or (messages, ...)"
+ "Eval function has no parameters. "
+ "Expected (solution_str, ...) or (messages, ...)"
)
first_param = params[0]
if first_param == "solution_str":
@@ -65,10 +66,10 @@ def _detect_mode(fn: Callable) -> str:
f"with 'osmosis eval' directly."
)
except (ValueError, TypeError) as e:
- raise EvalFnError(f"Cannot inspect eval function signature: {e}")
+ raise EvalFnError(f"Cannot inspect eval function signature: {e}") from e
@staticmethod
- def _extract_last_assistant_content(messages: List[Dict[str, Any]]) -> str:
+ def _extract_last_assistant_content(messages: list[dict[str, Any]]) -> str:
"""Extract content from the last assistant message."""
for msg in reversed(messages):
if msg.get("role") == "assistant":
@@ -79,9 +80,9 @@ def _extract_last_assistant_content(messages: List[Dict[str, Any]]) -> str:
async def __call__(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
ground_truth: str,
- metadata: Dict[str, Any],
+ metadata: dict[str, Any],
) -> float:
"""Call the eval function with normalized arguments.
@@ -148,15 +149,15 @@ def load_eval_fn(module_path: str) -> Callable:
try:
module = importlib.import_module(module_name)
except ImportError as e:
- raise EvalFnError(f"Cannot import module '{module_name}': {e}")
+ raise EvalFnError(f"Cannot import module '{module_name}': {e}") from e
try:
fn = getattr(module, attr_name)
- except AttributeError:
+ except AttributeError as e:
raise EvalFnError(
f"Module '{module_name}' has no attribute '{attr_name}'. "
f"Available: {[a for a in dir(module) if not a.startswith('_')]}"
- )
+ ) from e
if not callable(fn):
raise EvalFnError(f"'{attr_name}' in '{module_name}' is not callable")
@@ -164,7 +165,7 @@ def load_eval_fn(module_path: str) -> Callable:
return fn
-def load_eval_fns(module_paths: List[str]) -> List[EvalFnWrapper]:
+def load_eval_fns(module_paths: list[str]) -> list[EvalFnWrapper]:
"""Load multiple eval functions and wrap them.
Args:
@@ -189,6 +190,7 @@ def load_eval_fns(module_paths: List[str]) -> List[EvalFnWrapper]:
wrappers.append(wrapper)
return wrappers
+
__all__ = [
"EvalFnError",
"EvalFnWrapper",
diff --git a/osmosis_ai/rollout/eval/evaluation/report.py b/osmosis_ai/rollout/eval/evaluation/report.py
index 82de6826..86a9c471 100644
--- a/osmosis_ai/rollout/eval/evaluation/report.py
+++ b/osmosis_ai/rollout/eval/evaluation/report.py
@@ -8,8 +8,8 @@
from osmosis_ai.rollout.eval.common.cli import format_duration
if TYPE_CHECKING:
- from osmosis_ai.rollout.eval.evaluation.runner import EvalResult
from osmosis_ai.rollout.console import Console
+ from osmosis_ai.rollout.eval.evaluation.runner import EvalResult
def pass_at_k(n: int, c: int, k: int) -> float:
@@ -71,7 +71,7 @@ def _collect_pass_k_values(eval_summaries: dict[str, Any]) -> list[int]:
return pass_k_values
-def _build_model_data(result: "EvalResult", model: str) -> list[dict[str, Any]]:
+def _build_model_data(result: EvalResult, model: str) -> list[dict[str, Any]]:
"""Build per-model data list for unified table rendering.
In comparison mode, returns one entry per model from model_summaries.
@@ -83,9 +83,7 @@ def _build_model_data(result: "EvalResult", model: str) -> list[dict[str, Any]]:
avg_latency = (
ms.total_duration_ms / ms.total_runs if ms.total_runs > 0 else 0
)
- avg_tokens = (
- ms.total_tokens / ms.total_runs if ms.total_runs > 0 else 0
- )
+ avg_tokens = ms.total_tokens / ms.total_runs if ms.total_runs > 0 else 0
model_data.append(
{
"name": ms.model or ms.model_tag,
@@ -99,9 +97,7 @@ def _build_model_data(result: "EvalResult", model: str) -> list[dict[str, Any]]:
avg_latency = (
result.total_duration_ms / result.total_runs if result.total_runs > 0 else 0
)
- avg_tokens = (
- result.total_tokens / result.total_runs if result.total_runs > 0 else 0
- )
+ avg_tokens = result.total_tokens / result.total_runs if result.total_runs > 0 else 0
return [
{
"name": model,
@@ -112,9 +108,7 @@ def _build_model_data(result: "EvalResult", model: str) -> list[dict[str, Any]]:
]
-def format_eval_report(
- result: "EvalResult", console: "Console", model: str = ""
-) -> None:
+def format_eval_report(result: EvalResult, console: Console, model: str = "") -> None:
"""Print a formatted evaluation report to the console.
Renders a unified table with per-model rows for each eval function.
@@ -185,8 +179,8 @@ def _format_tables_rich(
rich_console: Any,
) -> None:
"""Render the performance and eval results tables using rich."""
- from rich.table import Table
from rich import box
+ from rich.table import Table
# --- Performance table ---
perf_table = Table(
@@ -227,8 +221,8 @@ def _format_tables_rich(
for k in pass_k_values:
table.add_column(f"pass@{k}", justify="right")
- for fn_idx, fn_name in enumerate(eval_fn_names):
- for m_idx, md in enumerate(model_data):
+ for _fn_idx, fn_name in enumerate(eval_fn_names):
+ for _m_idx, md in enumerate(model_data):
summary = md["eval_summaries"].get(fn_name)
if summary is None:
continue
@@ -262,7 +256,7 @@ def _format_tables_plain(
eval_fn_names: list[str],
model_data: list[dict[str, Any]],
pass_k_values: list[int],
- console: "Console",
+ console: Console,
) -> None:
"""Render the performance and eval results tables as plain text."""
# --- Performance table ---
diff --git a/osmosis_ai/rollout/eval/evaluation/runner.py b/osmosis_ai/rollout/eval/evaluation/runner.py
index 78dcd818..f0e9e52c 100644
--- a/osmosis_ai/rollout/eval/evaluation/runner.py
+++ b/osmosis_ai/rollout/eval/evaluation/runner.py
@@ -10,15 +10,16 @@
import inspect
import logging
import time
+from collections.abc import Callable
from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any
-from osmosis_ai.rollout.eval.evaluation.eval_fn import EvalFnWrapper
from osmosis_ai.rollout.core.base import RolloutAgentLoop
from osmosis_ai.rollout.eval.common.dataset import DatasetRow
-from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
+from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.common.runner import LocalRolloutRunner
+from osmosis_ai.rollout.eval.evaluation.eval_fn import EvalFnWrapper
logger = logging.getLogger(__name__)
@@ -54,11 +55,11 @@ class EvalRunResult:
run_index: int
success: bool
- scores: Dict[str, float] = field(default_factory=dict)
+ scores: dict[str, float] = field(default_factory=dict)
duration_ms: float = 0.0
tokens: int = 0
- error: Optional[str] = None
- model_tag: Optional[str] = None
+ error: str | None = None
+ model_tag: str | None = None
@dataclass
@@ -71,7 +72,7 @@ class EvalRowResult:
"""
row_index: int
- runs: List[EvalRunResult] = field(default_factory=list)
+ runs: list[EvalRunResult] = field(default_factory=list)
@dataclass
@@ -90,7 +91,7 @@ class EvalEvalSummary:
std: float = 0.0
min: float = 0.0
max: float = 0.0
- pass_at_k: Dict[int, float] = field(default_factory=dict)
+ pass_at_k: dict[int, float] = field(default_factory=dict)
@dataclass
@@ -108,7 +109,7 @@ class EvalModelSummary:
model: str
model_tag: str
- eval_summaries: Dict[str, EvalEvalSummary]
+ eval_summaries: dict[str, EvalEvalSummary]
total_runs: int
total_tokens: int
total_duration_ms: float
@@ -130,8 +131,8 @@ class EvalResult:
model_summaries: Per-model summaries (comparison mode only).
"""
- rows: List[EvalRowResult]
- eval_summaries: Dict[str, EvalEvalSummary]
+ rows: list[EvalRowResult]
+ eval_summaries: dict[str, EvalEvalSummary]
total_rows: int
total_runs: int
total_tokens: int
@@ -139,8 +140,8 @@ class EvalResult:
n_runs: int
pass_threshold: float
stopped_early: bool = False
- stop_reason: Optional[str] = None
- model_summaries: Optional[List[EvalModelSummary]] = None
+ stop_reason: str | None = None
+ model_summaries: list[EvalModelSummary] | None = None
class EvalRunner:
@@ -150,19 +151,21 @@ def __init__(
self,
agent_loop: RolloutAgentLoop,
llm_client: ExternalLLMClient,
- eval_fns: List[EvalFnWrapper],
+ eval_fns: list[EvalFnWrapper],
debug: bool = False,
- debug_dir: Optional[str] = None,
- llm_client_factory: Optional[Callable[[], Any]] = None,
- baseline_llm_client: Optional[ExternalLLMClient] = None,
- baseline_llm_client_factory: Optional[Callable[[], Any]] = None,
+ debug_dir: str | None = None,
+ llm_client_factory: Callable[[], Any] | None = None,
+ baseline_llm_client: ExternalLLMClient | None = None,
+ baseline_llm_client_factory: Callable[[], Any] | None = None,
) -> None:
self.agent_loop = agent_loop
self.llm_client = llm_client
self.eval_fns = eval_fns
self.debug = debug
self.debug_dir = debug_dir
- self._llm_client_factory = llm_client_factory or self._default_llm_client_factory
+ self._llm_client_factory = (
+ llm_client_factory or self._default_llm_client_factory
+ )
self._rollout_runner = LocalRolloutRunner(
agent_loop=agent_loop,
llm_client=llm_client,
@@ -177,7 +180,7 @@ def __init__(
self._baseline_llm_client_factory = (
baseline_llm_client_factory or self._default_baseline_llm_client_factory
)
- self._baseline_rollout_runner: Optional[LocalRolloutRunner] = None
+ self._baseline_rollout_runner: LocalRolloutRunner | None = None
if baseline_llm_client is not None:
self._baseline_rollout_runner = LocalRolloutRunner(
agent_loop=agent_loop,
@@ -243,9 +246,9 @@ async def run_single(
row_index: int,
run_index: int,
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
- runner: Optional[LocalRolloutRunner] = None,
- model_tag: Optional[str] = None,
+ completion_params: dict[str, Any] | None = None,
+ runner: LocalRolloutRunner | None = None,
+ model_tag: str | None = None,
) -> EvalRunResult:
"""Run agent once on a row and apply all eval functions.
@@ -289,8 +292,8 @@ async def run_single(
e,
fallback_started_at=run_start,
)
- setattr(e, "duration_ms", duration_ms)
- setattr(e, "tokens", tokens)
+ e.duration_ms = duration_ms
+ e.tokens = tokens
raise
if not test_result.success or test_result.result is None:
@@ -308,7 +311,7 @@ async def run_single(
ground_truth = row["ground_truth"]
metadata = dict(row)
- scores: Dict[str, float] = {}
+ scores: dict[str, float] = {}
for eval_fn in self.eval_fns:
try:
score = await eval_fn(messages, ground_truth, metadata)
@@ -316,7 +319,10 @@ async def run_single(
except Exception as e:
logger.warning(
"Eval function '%s' failed on row %d run %d: %s",
- eval_fn.name, row_index, run_index, e,
+ eval_fn.name,
+ row_index,
+ run_index,
+ e,
)
scores[eval_fn.name] = 0.0
@@ -331,12 +337,12 @@ async def run_single(
async def run_eval(
self,
- rows: List[DatasetRow],
+ rows: list[DatasetRow],
n_runs: int = 1,
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
+ completion_params: dict[str, Any] | None = None,
pass_threshold: float = 1.0,
- on_progress: Optional[Callable[[int, int, EvalRunResult], None]] = None,
+ on_progress: Callable[[int, int, EvalRunResult], None] | None = None,
start_index: int = 0,
batch_size: int = 1,
) -> EvalResult:
@@ -371,14 +377,12 @@ async def run_eval(
)
total_start = time.monotonic()
- row_results: List[EvalRowResult] = []
- model_tags = (
- ["primary", "baseline"] if self.has_baseline else [None]
- )
+ row_results: list[EvalRowResult] = []
+ model_tags = ["primary", "baseline"] if self.has_baseline else [None]
total = len(rows) * n_runs * len(model_tags)
current = 0
stopped_early = False
- stop_reason: Optional[str] = None
+ stop_reason: str | None = None
for i, row in enumerate(rows):
row_index = start_index + i
@@ -432,12 +436,10 @@ async def run_eval(
break
total_duration_ms = (time.monotonic() - total_start) * 1000
- total_tokens = sum(
- run.tokens for row in row_results for run in row.runs
- )
+ total_tokens = sum(run.tokens for row in row_results for run in row.runs)
# Compute per-model summaries if baseline is configured
- model_summaries: Optional[List[EvalModelSummary]] = None
+ model_summaries: list[EvalModelSummary] | None = None
if self.has_baseline:
model_summaries = self._compute_model_summaries(
row_results, n_runs, pass_threshold
@@ -471,12 +473,12 @@ async def run_eval(
async def _run_eval_concurrent(
self,
- rows: List[DatasetRow],
+ rows: list[DatasetRow],
n_runs: int,
max_turns: int,
- completion_params: Optional[Dict[str, Any]],
+ completion_params: dict[str, Any] | None,
pass_threshold: float,
- on_progress: Optional[Callable[[int, int, EvalRunResult], None]],
+ on_progress: Callable[[int, int, EvalRunResult], None] | None,
start_index: int,
batch_size: int,
) -> EvalResult:
@@ -488,7 +490,7 @@ async def _run_eval_concurrent(
are created and the batch_size is split between them.
"""
total_start = time.monotonic()
- model_tags: List[Optional[str]] = (
+ model_tags: list[str | None] = (
["primary", "baseline"] if self.has_baseline else [None]
)
total = len(rows) * n_runs * len(model_tags)
@@ -505,15 +507,20 @@ async def _run_eval_concurrent(
primary_pool_size = pool_size
baseline_pool_size = 0
- primary_runners = [self._create_rollout_runner() for _ in range(primary_pool_size)]
+ primary_runners = [
+ self._create_rollout_runner() for _ in range(primary_pool_size)
+ ]
primary_pool: asyncio.Queue[LocalRolloutRunner] = asyncio.Queue()
for runner in primary_runners:
primary_pool.put_nowait(runner)
- baseline_runners: List[LocalRolloutRunner] = []
+ baseline_runners: list[LocalRolloutRunner] = []
baseline_pool: asyncio.Queue[LocalRolloutRunner] = asyncio.Queue()
if baseline_pool_size > 0:
- baseline_runners = [self._create_baseline_rollout_runner() for _ in range(baseline_pool_size)]
+ baseline_runners = [
+ self._create_baseline_rollout_runner()
+ for _ in range(baseline_pool_size)
+ ]
for runner in baseline_runners:
baseline_pool.put_nowait(runner)
@@ -521,8 +528,8 @@ async def _run_eval_concurrent(
completed = 0
stopped_early = False
- stop_reason: Optional[str] = None
- completed_results: List[Tuple[int, EvalRunResult]] = []
+ stop_reason: str | None = None
+ completed_results: list[tuple[int, EvalRunResult]] = []
async def _close_runner_pool() -> None:
"""Best-effort cleanup for per-runner LLM clients."""
@@ -535,14 +542,18 @@ async def _close_runner_pool() -> None:
if inspect.isawaitable(maybe_awaitable):
await maybe_awaitable
except Exception:
- logger.debug("Failed to close concurrent runner client", exc_info=True)
+ logger.debug(
+ "Failed to close concurrent runner client", exc_info=True
+ )
async def _run_one(
- row: DatasetRow, row_index: int, run_index: int,
- model_tag: Optional[str],
- ) -> Tuple[int, EvalRunResult]:
+ row: DatasetRow,
+ row_index: int,
+ run_index: int,
+ model_tag: str | None,
+ ) -> tuple[int, EvalRunResult]:
nonlocal completed
- runner: Optional[LocalRolloutRunner] = None
+ runner: LocalRolloutRunner | None = None
run_start = time.monotonic()
target_pool = baseline_pool if model_tag == "baseline" else primary_pool
try:
@@ -587,7 +598,7 @@ async def _run_one(
target_pool.put_nowait(runner)
# Build work items: interleave primary/baseline per (row, run_idx)
- work_items: List[Tuple[DatasetRow, int, int, Optional[str]]] = []
+ work_items: list[tuple[DatasetRow, int, int, str | None]] = []
for i, row in enumerate(rows):
row_index = start_index + i
for run_idx in range(n_runs):
@@ -597,18 +608,18 @@ async def _run_one(
try:
cursor = 0
while cursor < len(work_items):
- batch = work_items[cursor:cursor + batch_size]
- tasks: List[asyncio.Task[Tuple[int, EvalRunResult]]] = [
+ batch = work_items[cursor : cursor + batch_size]
+ tasks: list[asyncio.Task[tuple[int, EvalRunResult]]] = [
asyncio.create_task(_run_one(row, row_index, run_idx, tag))
for row, row_index, run_idx, tag in batch
]
- systemic_error: Optional[str] = None
+ systemic_error: str | None = None
try:
for task in asyncio.as_completed(tasks):
try:
- _row_index, result = await task
+ _row_index, _result = await task
except SystemicProviderError as e:
# Record error but keep awaiting remaining
# tasks in this batch so their results are collected.
@@ -633,7 +644,7 @@ async def _run_one(
await _close_runner_pool()
# Organise flat results back into per-row structure.
- row_results_map: Dict[int, EvalRowResult] = {}
+ row_results_map: dict[int, EvalRowResult] = {}
for row_index, run_result in completed_results:
if row_index not in row_results_map:
row_results_map[row_index] = EvalRowResult(row_index=row_index)
@@ -650,12 +661,10 @@ async def _run_one(
row_results = [row_results_map[i] for i in sorted(row_results_map.keys())]
total_duration_ms = (time.monotonic() - total_start) * 1000
- total_tokens = sum(
- run.tokens for row in row_results for run in row.runs
- )
+ total_tokens = sum(run.tokens for row in row_results for run in row.runs)
# Compute per-model summaries if baseline is configured
- model_summaries: Optional[List[EvalModelSummary]] = None
+ model_summaries: list[EvalModelSummary] | None = None
if self.has_baseline:
model_summaries = self._compute_model_summaries(
row_results, n_runs, pass_threshold
@@ -689,10 +698,10 @@ async def _run_one(
def _compute_summaries(
self,
- row_results: List[EvalRowResult],
+ row_results: list[EvalRowResult],
n_runs: int,
pass_threshold: float,
- ) -> Dict[str, EvalEvalSummary]:
+ ) -> dict[str, EvalEvalSummary]:
"""Compute per-eval-function summary statistics.
Failed runs and missing eval scores are treated as 0.0 so reliability
@@ -700,13 +709,13 @@ def _compute_summaries(
"""
import math
- summaries: Dict[str, EvalEvalSummary] = {}
+ summaries: dict[str, EvalEvalSummary] = {}
for eval_fn in self.eval_fns:
name = eval_fn.name
# Collect all scores for this eval fn. Missing scores count as 0.
- all_scores: List[float] = []
+ all_scores: list[float] = []
for row in row_results:
for run in row.runs:
all_scores.append(run.scores.get(name, 0.0))
@@ -736,10 +745,11 @@ def _compute_summaries(
# Compute pass@k per row, then average.
# Use n_runs (not len(row.runs)) so that missing runs
# (e.g. from early stopping) are treated as failures.
- row_pass_at_k: List[float] = []
+ row_pass_at_k: list[float] = []
for row in row_results:
c = sum(
- 1 for run in row.runs
+ 1
+ for run in row.runs
if run.scores.get(name, 0.0) >= pass_threshold
)
n = max(len(row.runs), n_runs)
@@ -754,28 +764,30 @@ def _compute_summaries(
def _filter_runs_by_tag(
self,
- row_results: List[EvalRowResult],
+ row_results: list[EvalRowResult],
model_tag: str,
- ) -> List[EvalRowResult]:
+ ) -> list[EvalRowResult]:
"""Build row results containing only runs matching the given model_tag."""
- filtered: List[EvalRowResult] = []
+ filtered: list[EvalRowResult] = []
for row in row_results:
tagged_runs = [r for r in row.runs if r.model_tag == model_tag]
if tagged_runs:
- filtered.append(EvalRowResult(
- row_index=row.row_index,
- runs=tagged_runs,
- ))
+ filtered.append(
+ EvalRowResult(
+ row_index=row.row_index,
+ runs=tagged_runs,
+ )
+ )
return filtered
def _compute_model_summaries(
self,
- row_results: List[EvalRowResult],
+ row_results: list[EvalRowResult],
n_runs: int,
pass_threshold: float,
- ) -> List[EvalModelSummary]:
+ ) -> list[EvalModelSummary]:
"""Compute per-model summary statistics for comparison mode."""
- summaries: List[EvalModelSummary] = []
+ summaries: list[EvalModelSummary] = []
for tag in ("primary", "baseline"):
filtered = self._filter_runs_by_tag(row_results, tag)
eval_summaries = self._compute_summaries(filtered, n_runs, pass_threshold)
@@ -785,16 +797,19 @@ def _compute_model_summaries(
model_name = self.llm_client.display_name
elif self._baseline_llm_client is not None:
model_name = self._baseline_llm_client.display_name
- summaries.append(EvalModelSummary(
- model=model_name,
- model_tag=tag,
- eval_summaries=eval_summaries,
- total_runs=len(all_runs),
- total_tokens=sum(r.tokens for r in all_runs),
- total_duration_ms=sum(r.duration_ms for r in all_runs),
- ))
+ summaries.append(
+ EvalModelSummary(
+ model=model_name,
+ model_tag=tag,
+ eval_summaries=eval_summaries,
+ total_runs=len(all_runs),
+ total_tokens=sum(r.tokens for r in all_runs),
+ total_duration_ms=sum(r.duration_ms for r in all_runs),
+ )
+ )
return summaries
+
__all__ = [
"EvalEvalSummary",
"EvalModelSummary",
diff --git a/osmosis_ai/rollout/eval/test_mode/__init__.py b/osmosis_ai/rollout/eval/test_mode/__init__.py
index 18774dab..032397d6 100644
--- a/osmosis_ai/rollout/eval/test_mode/__init__.py
+++ b/osmosis_ai/rollout/eval/test_mode/__init__.py
@@ -1,7 +1,7 @@
"""Test mode for RolloutAgentLoop validation."""
-from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.common.errors import ProviderError
+from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.test_mode.interactive import (
InteractiveLLMClient,
InteractiveRunner,
@@ -9,17 +9,17 @@
)
from osmosis_ai.rollout.eval.test_mode.runner import (
LocalTestBatchResult,
- LocalTestRunResult,
LocalTestRunner,
+ LocalTestRunResult,
)
__all__ = [
- "LocalTestRunner",
- "LocalTestRunResult",
- "LocalTestBatchResult",
- "InteractiveRunner",
+ "ExternalLLMClient",
"InteractiveLLMClient",
+ "InteractiveRunner",
"InteractiveStep",
- "ExternalLLMClient",
+ "LocalTestBatchResult",
+ "LocalTestRunResult",
+ "LocalTestRunner",
"ProviderError",
]
diff --git a/osmosis_ai/rollout/eval/test_mode/cli.py b/osmosis_ai/rollout/eval/test_mode/cli.py
index 90bab5e8..061dcedc 100644
--- a/osmosis_ai/rollout/eval/test_mode/cli.py
+++ b/osmosis_ai/rollout/eval/test_mode/cli.py
@@ -11,7 +11,7 @@
import logging
from dataclasses import dataclass
from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.console import Console
from osmosis_ai.rollout.eval.common.cli import (
@@ -29,7 +29,10 @@
from osmosis_ai.rollout.core.base import RolloutAgentLoop
from osmosis_ai.rollout.eval.common.dataset import DatasetRow
from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
- from osmosis_ai.rollout.eval.test_mode.runner import LocalTestBatchResult, LocalTestRunResult
+ from osmosis_ai.rollout.eval.test_mode.runner import (
+ LocalTestBatchResult,
+ LocalTestRunResult,
+ )
logger = logging.getLogger(__name__)
@@ -39,10 +42,10 @@
class _SetupResult:
"""Result of setup phase containing initialized components."""
- agent_loop: "RolloutAgentLoop"
- llm_client: "ExternalLLMClient"
- rows: List["DatasetRow"]
- completion_params: Dict[str, Any]
+ agent_loop: RolloutAgentLoop
+ llm_client: ExternalLLMClient
+ rows: list[DatasetRow]
+ completion_params: dict[str, Any]
class TestCommand:
@@ -200,7 +203,7 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None:
def run(self, args: argparse.Namespace) -> int:
return asyncio.run(self._run_async(args))
- def _validate_args(self, args: argparse.Namespace) -> Optional[str]:
+ def _validate_args(self, args: argparse.Namespace) -> str | None:
if args.module and args.mcp:
return "--module and --mcp are mutually exclusive."
if not args.module and not args.mcp:
@@ -216,7 +219,9 @@ def _print_header(self, args: argparse.Namespace) -> None:
from osmosis_ai.consts import PACKAGE_VERSION
mode_suffix = " (Interactive Mode)" if args.interactive else ""
- self.console.print(f"osmosis-rollout-test v{PACKAGE_VERSION}{mode_suffix}", style="bold")
+ self.console.print(
+ f"osmosis-rollout-test v{PACKAGE_VERSION}{mode_suffix}", style="bold"
+ )
self.console.print()
async def _run_interactive_mode(
@@ -238,7 +243,9 @@ async def _run_interactive_mode(
await interactive_runner.run_interactive_session(
rows=setup.rows,
max_turns=args.max_turns,
- completion_params=setup.completion_params if setup.completion_params else None,
+ completion_params=setup.completion_params
+ if setup.completion_params
+ else None,
initial_row=args.row,
row_offset=args.offset,
)
@@ -256,7 +263,7 @@ async def _run_batch_mode(
self,
args: argparse.Namespace,
setup: _SetupResult,
- ) -> "LocalTestBatchResult":
+ ) -> LocalTestBatchResult:
from osmosis_ai.rollout.eval.test_mode.runner import LocalTestRunner
runner = LocalTestRunner(
@@ -265,9 +272,7 @@ async def _run_batch_mode(
debug=args.debug,
)
- def on_progress(
- current: int, total: int, result: "LocalTestRunResult"
- ) -> None:
+ def on_progress(current: int, total: int, result: LocalTestRunResult) -> None:
if args.quiet:
return
@@ -279,7 +284,9 @@ def on_progress(
error_suffix = ""
if not result.success and result.error:
error_text = result.error.replace("\n", " ")
- error_msg = error_text[:47] + "..." if len(error_text) > 50 else error_text
+ error_msg = (
+ error_text[:47] + "..." if len(error_text) > 50 else error_text
+ )
error_suffix = f" - {error_msg}"
status_styled = self.console.format_styled(status, status_style)
@@ -296,14 +303,16 @@ def on_progress(
batch_result = await runner.run_batch(
rows=setup.rows,
max_turns=args.max_turns,
- completion_params=setup.completion_params if setup.completion_params else None,
+ completion_params=setup.completion_params
+ if setup.completion_params
+ else None,
on_progress=on_progress,
start_index=args.offset,
)
return batch_result
- def _print_summary(self, batch_result: "LocalTestBatchResult") -> None:
+ def _print_summary(self, batch_result: LocalTestBatchResult) -> None:
self.console.print()
self.console.print("Summary:", style="bold")
self.console.print(f" Total: {batch_result.total}")
@@ -324,18 +333,26 @@ def _print_summary(self, batch_result: "LocalTestBatchResult") -> None:
self.console.print(f" Passed: {passed_text}")
self.console.print(f" Failed: {failed_text}")
- self.console.print(f" Duration: {format_duration(batch_result.total_duration_ms)}")
- self.console.print(f" Total tokens: {format_tokens(batch_result.total_tokens)}")
+ self.console.print(
+ f" Duration: {format_duration(batch_result.total_duration_ms)}"
+ )
+ self.console.print(
+ f" Total tokens: {format_tokens(batch_result.total_tokens)}"
+ )
if batch_result.stopped_early:
- reason = f" Reason: {batch_result.stop_reason}" if batch_result.stop_reason else ""
+ reason = (
+ f" Reason: {batch_result.stop_reason}"
+ if batch_result.stop_reason
+ else ""
+ )
self.console.print(
f" Stopped early due to systemic provider error.{reason}",
style="red",
)
def _write_output(
- self, args: argparse.Namespace, batch_result: "LocalTestBatchResult"
+ self, args: argparse.Namespace, batch_result: LocalTestBatchResult
) -> None:
if not args.output:
return
@@ -463,4 +480,5 @@ async def _run_async(self, args: argparse.Namespace) -> int:
self._write_output(args, batch_result)
return 1 if batch_result.failed > 0 else 0
+
__all__ = ["TestCommand"]
diff --git a/osmosis_ai/rollout/eval/test_mode/interactive.py b/osmosis_ai/rollout/eval/test_mode/interactive.py
index 1130a661..03eb5eaf 100644
--- a/osmosis_ai/rollout/eval/test_mode/interactive.py
+++ b/osmosis_ai/rollout/eval/test_mode/interactive.py
@@ -6,12 +6,12 @@
from __future__ import annotations
-import asyncio
import copy
import json
import time
+from collections.abc import Callable
from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any
from osmosis_ai.rollout.client import CompletionsResult
from osmosis_ai.rollout.console import Console
@@ -22,8 +22,8 @@
)
from osmosis_ai.rollout.core.schemas import OpenAIFunctionToolSchema, RolloutMetrics
from osmosis_ai.rollout.eval.common.dataset import DatasetRow, dataset_row_to_request
-from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.common.errors import ToolValidationError
+from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
from osmosis_ai.rollout.eval.common.runner import validate_tools
@@ -41,9 +41,9 @@ class InteractiveStep:
turn: int
step_type: str # "llm_response", "tool_result", "final"
- message: Optional[Dict[str, Any]] = None
- tool_calls: Optional[List[Dict[str, Any]]] = None
- finish_reason: Optional[str] = None
+ message: dict[str, Any] | None = None
+ tool_calls: list[dict[str, Any]] | None = None
+ finish_reason: str | None = None
class InteractiveLLMClient:
@@ -57,7 +57,7 @@ def __init__(
self,
wrapped_client: ExternalLLMClient,
on_step: Callable[[InteractiveStep], bool],
- on_messages_updated: Optional[Callable[[List[Dict[str, Any]]], None]] = None,
+ on_messages_updated: Callable[[list[dict[str, Any]]], None] | None = None,
) -> None:
"""Initialize the interactive client wrapper.
@@ -75,7 +75,7 @@ def __init__(
self._turn = 0
self._aborted = False
- def set_tools(self, tools: List[Any]) -> None:
+ def set_tools(self, tools: list[Any]) -> None:
"""Delegate to wrapped client."""
self._wrapped.set_tools(tools)
@@ -95,7 +95,7 @@ def get_metrics(self) -> RolloutMetrics:
async def chat_completions(
self,
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
**kwargs: Any,
) -> CompletionsResult:
"""Make a chat completion and pause for user interaction.
@@ -148,7 +148,7 @@ async def close(self) -> None:
"""Delegate to wrapped client."""
await self._wrapped.close()
- async def __aenter__(self) -> "InteractiveLLMClient":
+ async def __aenter__(self) -> InteractiveLLMClient:
"""Async context manager entry."""
await self._wrapped.__aenter__()
return self
@@ -165,7 +165,7 @@ class InteractiveRunner:
"""
# Command definitions: maps aliases to (handler_method_name, description)
- _COMMANDS: Dict[str, Tuple[str, str]] = {
+ _COMMANDS: dict[str, tuple[str, str]] = {
"n": ("_cmd_next", "Continue to next step"),
"next": ("_cmd_next", "Continue to next step"),
"": ("_cmd_next", "Continue to next step"),
@@ -181,7 +181,7 @@ class InteractiveRunner:
}
# Short aliases for help display (primary commands only)
- _HELP_COMMANDS: List[Tuple[str, str, str]] = [
+ _HELP_COMMANDS: list[tuple[str, str, str]] = [
("n/next", "", "Continue to next step"),
("c", "", "Continue to completion (no more pauses)"),
("m", "", "Show all messages"),
@@ -208,8 +208,8 @@ def __init__(
self.console = Console()
# State for interactive session
- self._current_messages: List[Dict[str, Any]] = []
- self._current_tools: List[OpenAIFunctionToolSchema] = []
+ self._current_messages: list[dict[str, Any]] = []
+ self._current_tools: list[OpenAIFunctionToolSchema] = []
self._auto_continue = False
self._last_msg_count_seen = 0
@@ -217,9 +217,9 @@ def _print_separator(self, title: str = "") -> None:
"""Print a separator line."""
self.console.separator(title)
- def _build_tool_name_map(self, messages: List[Dict[str, Any]]) -> Dict[str, str]:
+ def _build_tool_name_map(self, messages: list[dict[str, Any]]) -> dict[str, str]:
"""Build mapping from tool_call_id to function name."""
- name_map: Dict[str, str] = {}
+ name_map: dict[str, str] = {}
for msg in messages:
if msg.get("role") == "assistant":
for tc in msg.get("tool_calls") or []:
@@ -231,9 +231,9 @@ def _build_tool_name_map(self, messages: List[Dict[str, Any]]) -> Dict[str, str]
def _print_message(
self,
- msg: Dict[str, Any],
+ msg: dict[str, Any],
prefix: str = "",
- tool_name_map: Optional[Dict[str, str]] = None,
+ tool_name_map: dict[str, str] | None = None,
) -> None:
"""Print a message in a formatted way."""
role = msg.get("role", "unknown")
@@ -258,13 +258,17 @@ def _print_message(
}
style, label = role_styles.get(role, (None, role.capitalize()))
- styled_label = self.console.format_styled(f"[{label}]", style) if style else f"[{label}]"
+ styled_label = (
+ self.console.format_styled(f"[{label}]", style) if style else f"[{label}]"
+ )
self.console.print(f"{prefix}{styled_label} {self.console.escape(content)}")
# Print tool calls if present
tool_calls = msg.get("tool_calls", [])
if tool_calls:
- self.console.print(f"\n{prefix}{self.console.format_styled('Tool calls:', 'yellow')}")
+ self.console.print(
+ f"\n{prefix}{self.console.format_styled('Tool calls:', 'yellow')}"
+ )
for tc in tool_calls:
func = tc.get("function", {})
name = func.get("name", "unknown")
@@ -279,14 +283,18 @@ def _print_message(
# Indent multi-line args
args_lines = args_str.split("\n")
if len(args_lines) > 1:
- args_display = args_lines[0] + "\n" + "\n".join(
- " " + line for line in args_lines[1:]
+ args_display = (
+ args_lines[0]
+ + "\n"
+ + "\n".join(" " + line for line in args_lines[1:])
)
else:
args_display = args_str
styled_name = self.console.format_styled(name, "cyan")
- self.console.print(f"{prefix} • {styled_name}({self.console.escape(args_display)})")
+ self.console.print(
+ f"{prefix} • {styled_name}({self.console.escape(args_display)})"
+ )
def _print_step(self, step: InteractiveStep) -> None:
"""Print a step in the interactive session."""
@@ -301,7 +309,7 @@ def _print_step(self, step: InteractiveStep) -> None:
self.console.print()
- def _print_initial_messages(self, messages: List[Dict[str, Any]]) -> None:
+ def _print_initial_messages(self, messages: list[dict[str, Any]]) -> None:
"""Print the initial conversation messages."""
self._print_separator("Initial Messages")
self.console.print()
@@ -309,7 +317,7 @@ def _print_initial_messages(self, messages: List[Dict[str, Any]]) -> None:
self._print_message(msg)
self.console.print()
- def _print_tools(self, tools: List[OpenAIFunctionToolSchema]) -> None:
+ def _print_tools(self, tools: list[OpenAIFunctionToolSchema]) -> None:
"""Print available tools."""
self._print_separator("Available Tools")
self.console.print()
@@ -324,7 +332,7 @@ def _print_tools(self, tools: List[OpenAIFunctionToolSchema]) -> None:
self.console.print(f" Parameters: {', '.join(props.keys())}")
self.console.print()
- def _print_all_messages(self, messages: List[Dict[str, Any]]) -> None:
+ def _print_all_messages(self, messages: list[dict[str, Any]]) -> None:
"""Print all messages in the conversation."""
self._print_separator("All Messages")
self.console.print()
@@ -336,10 +344,10 @@ def _print_all_messages(self, messages: List[Dict[str, Any]]) -> None:
def _print_result(
self,
- result: Optional[RolloutResult],
+ result: RolloutResult | None,
duration_ms: float,
metrics: RolloutMetrics,
- error: Optional[str] = None,
+ error: str | None = None,
) -> None:
"""Print the final result."""
self._print_separator("Result")
@@ -358,9 +366,11 @@ def _print_result(
if result.finish_reason:
self.console.print(f"Finish reason: {result.finish_reason}")
- self.console.print(f"Duration: {duration_ms/1000:.2f}s")
+ self.console.print(f"Duration: {duration_ms / 1000:.2f}s")
total_tokens = metrics.prompt_tokens + metrics.response_tokens
- self.console.print(f"Tokens: {total_tokens:,} ({metrics.num_llm_calls} LLM calls)")
+ self.console.print(
+ f"Tokens: {total_tokens:,} ({metrics.num_llm_calls} LLM calls)"
+ )
self.console.print()
def _post_result_prompt(self) -> None:
@@ -408,27 +418,27 @@ def _get_user_input(self) -> str:
# If should_return is True, _handle_step returns return_value
# If should_return is False, the command loop continues
- def _cmd_next(self) -> Tuple[bool, bool]:
+ def _cmd_next(self) -> tuple[bool, bool]:
"""Continue to next step."""
return (True, True)
- def _cmd_continue(self) -> Tuple[bool, bool]:
+ def _cmd_continue(self) -> tuple[bool, bool]:
"""Auto-continue to completion."""
self._auto_continue = True
self.console.print("Continuing to completion...", style="dim")
return (True, True)
- def _cmd_messages(self) -> Tuple[bool, bool]:
+ def _cmd_messages(self) -> tuple[bool, bool]:
"""Show all messages."""
self._print_all_messages(self._current_messages)
return (False, False)
- def _cmd_tools(self) -> Tuple[bool, bool]:
+ def _cmd_tools(self) -> tuple[bool, bool]:
"""Show tools."""
self._print_tools(self._current_tools)
return (False, False)
- def _cmd_quit(self) -> Tuple[bool, bool]:
+ def _cmd_quit(self) -> tuple[bool, bool]:
"""Abort execution."""
self.console.print("Aborting execution...", style="yellow")
return (True, False)
@@ -438,7 +448,7 @@ def _style_command_hint(self, command: str) -> str:
# format_styled already calls rich_escape internally, so pass raw text.
return self.console.format_styled(f"[{command}]", "cyan")
- def _print_help(self, unknown_cmd: Optional[str] = None) -> None:
+ def _print_help(self, unknown_cmd: str | None = None) -> None:
"""Print help for available commands."""
if unknown_cmd is not None:
self.console.print(f"Unknown command: {unknown_cmd}", style="red")
@@ -473,8 +483,8 @@ async def run_single_interactive(
row: DatasetRow,
row_index: int,
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
- ) -> Tuple[Optional[RolloutResult], Optional[str]]:
+ completion_params: dict[str, Any] | None = None,
+ ) -> tuple[RolloutResult | None, str | None]:
"""Run a single row interactively.
Args:
@@ -524,9 +534,9 @@ async def run_single_interactive(
self._print_tools(tools)
# Create callback to update current messages for the 'm' command
- def update_messages(messages: List[Dict[str, Any]]) -> None:
+ def update_messages(messages: list[dict[str, Any]]) -> None:
# Print new tool messages since last update
- new_msgs = messages[self._last_msg_count_seen:]
+ new_msgs = messages[self._last_msg_count_seen :]
if any(m.get("role") == "tool" for m in new_msgs):
tool_name_map = self._build_tool_name_map(messages)
self.console.print()
@@ -599,10 +609,10 @@ def update_messages(messages: List[Dict[str, Any]]) -> None:
async def run_interactive_session(
self,
- rows: List[DatasetRow],
+ rows: list[DatasetRow],
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
- initial_row: Optional[int] = None,
+ completion_params: dict[str, Any] | None = None,
+ initial_row: int | None = None,
row_offset: int = 0,
) -> None:
"""Run an interactive testing session.
@@ -652,7 +662,9 @@ async def run_interactive_session(
self.console.print()
range_styled = self.console.format_styled(f"{start_row}-{end_row}", "cyan")
q_styled = self.console.format_styled("q", "cyan")
- self.console.print(f"Select a row to test [{range_styled}] or '{q_styled}' to quit:")
+ self.console.print(
+ f"Select a row to test [{range_styled}] or '{q_styled}' to quit:"
+ )
try:
user_input = self.console.input("> ", style="bold").strip().lower()
@@ -690,6 +702,7 @@ async def run_interactive_session(
self.console.print("\nInteractive session ended.", style="dim")
+
__all__ = [
"InteractiveLLMClient",
"InteractiveRunner",
diff --git a/osmosis_ai/rollout/eval/test_mode/runner.py b/osmosis_ai/rollout/eval/test_mode/runner.py
index 688fadde..80286520 100644
--- a/osmosis_ai/rollout/eval/test_mode/runner.py
+++ b/osmosis_ai/rollout/eval/test_mode/runner.py
@@ -6,14 +6,18 @@
from __future__ import annotations
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.eval.common.runner import (
LocalBatchResult as LocalTestBatchResult,
+)
+from osmosis_ai.rollout.eval.common.runner import (
LocalRolloutRunner,
- LocalRunResult as LocalTestRunResult,
validate_tools,
)
+from osmosis_ai.rollout.eval.common.runner import (
+ LocalRunResult as LocalTestRunResult,
+)
if TYPE_CHECKING:
from osmosis_ai.rollout.core.base import RolloutAgentLoop
@@ -26,16 +30,18 @@ class LocalTestRunner(LocalRolloutRunner):
def __init__(
self,
- agent_loop: "RolloutAgentLoop",
- llm_client: "ExternalLLMClient",
+ agent_loop: RolloutAgentLoop,
+ llm_client: ExternalLLMClient,
debug: bool = False,
- debug_dir: Optional[str] = None,
+ debug_dir: str | None = None,
) -> None:
super().__init__(
agent_loop=agent_loop,
llm_client=llm_client,
debug=debug,
- debug_dir=debug_dir if debug_dir is not None else ("./test_debug" if debug else None),
+ debug_dir=debug_dir
+ if debug_dir is not None
+ else ("./test_debug" if debug else None),
rollout_id_prefix="test",
request_metadata={
"execution_mode": "test",
@@ -45,10 +51,10 @@ def __init__(
async def run_single(
self,
- row: "DatasetRow",
+ row: DatasetRow,
row_index: int,
max_turns: int = 10,
- completion_params: Optional[Dict[str, Any]] = None,
+ completion_params: dict[str, Any] | None = None,
) -> LocalTestRunResult:
return await super().run_single(
row=row,
@@ -58,4 +64,10 @@ async def run_single(
rollout_id=f"test-{row_index}",
)
-__all__ = ["LocalTestBatchResult", "LocalTestRunResult", "LocalTestRunner", "validate_tools"]
+
+__all__ = [
+ "LocalTestBatchResult",
+ "LocalTestRunResult",
+ "LocalTestRunner",
+ "validate_tools",
+]
diff --git a/osmosis_ai/rollout/mcp/agent_loop.py b/osmosis_ai/rollout/mcp/agent_loop.py
index 7a08ea42..d1d16a95 100644
--- a/osmosis_ai/rollout/mcp/agent_loop.py
+++ b/osmosis_ai/rollout/mcp/agent_loop.py
@@ -8,7 +8,7 @@
import logging
import time
-from typing import Any, Dict, List, Optional
+from typing import Any
from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext, RolloutResult
from osmosis_ai.rollout.core.schemas import (
@@ -32,7 +32,8 @@
# MCP Tool → OpenAI Schema conversion
# ---------------------------------------------------------------------------
-def _resolve_ref(ref: str, defs: Dict[str, Any]) -> Dict[str, Any]:
+
+def _resolve_ref(ref: str, defs: dict[str, Any]) -> dict[str, Any]:
"""Resolve a ``$ref`` pointer like ``#/$defs/Color`` against *defs*."""
parts = ref.lstrip("#/").split("/")
node: Any = {"$defs": defs}
@@ -42,8 +43,8 @@ def _resolve_ref(ref: str, defs: Dict[str, Any]) -> Dict[str, Any]:
def _json_schema_to_property(
- schema: Dict[str, Any],
- defs: Dict[str, Any],
+ schema: dict[str, Any],
+ defs: dict[str, Any],
) -> OpenAIFunctionPropertySchema:
"""Convert a single JSON Schema property to an OpenAIFunctionPropertySchema.
@@ -64,7 +65,9 @@ def _json_schema_to_property(
prop_type_raw = schema.get("type", "string")
if isinstance(prop_type_raw, list):
# JSON Schema allows union types like ["integer", "null"].
- non_null_types = [t for t in prop_type_raw if isinstance(t, str) and t != "null"]
+ non_null_types = [
+ t for t in prop_type_raw if isinstance(t, str) and t != "null"
+ ]
prop_type = non_null_types[0] if non_null_types else "string"
elif isinstance(prop_type_raw, str):
prop_type = prop_type_raw
@@ -88,14 +91,14 @@ def _json_schema_to_property(
def _mcp_tool_to_openai(
name: str,
description: str,
- parameters: Dict[str, Any],
+ parameters: dict[str, Any],
) -> OpenAIFunctionToolSchema:
"""Convert an MCP FunctionTool's metadata to an OpenAI tool schema."""
defs = parameters.get("$defs", {})
raw_props = parameters.get("properties", {})
required = parameters.get("required", [])
- properties: Dict[str, OpenAIFunctionPropertySchema] = {}
+ properties: dict[str, OpenAIFunctionPropertySchema] = {}
for prop_name, prop_schema in raw_props.items():
properties[prop_name] = _json_schema_to_property(prop_schema, defs)
@@ -117,6 +120,7 @@ def _mcp_tool_to_openai(
# MCPAgentLoop
# ---------------------------------------------------------------------------
+
class MCPAgentLoop(RolloutAgentLoop):
"""Agent loop that executes MCP tools from a FastMCP server instance.
@@ -130,10 +134,10 @@ class MCPAgentLoop(RolloutAgentLoop):
def __init__(self, mcp_server: Any, agent_name: str = "mcp_agent") -> None:
self.name = agent_name # type: ignore[misc]
self._tool_manager = mcp_server._tool_manager
- self._openai_schemas: List[OpenAIFunctionToolSchema] = []
+ self._openai_schemas: list[OpenAIFunctionToolSchema] = []
# Build cached OpenAI schemas from registered MCP tools
- for tool_name, tool in self._tool_manager._tools.items():
+ for _tool_name, tool in self._tool_manager._tools.items():
self._openai_schemas.append(
_mcp_tool_to_openai(
name=tool.name,
@@ -142,13 +146,13 @@ def __init__(self, mcp_server: Any, agent_name: str = "mcp_agent") -> None:
)
)
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return list(self._openai_schemas)
async def run(self, ctx: RolloutContext) -> RolloutResult:
messages = list(ctx.request.messages)
- for turn in range(ctx.request.max_turns):
+ for _turn in range(ctx.request.max_turns):
result = await ctx.chat(messages, **ctx.request.completion_params)
messages.append(result.message)
@@ -164,7 +168,9 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
# Extract value: prefer structured_content, fallback to content text
value = _extract_tool_value(tool_result)
- messages.append(create_tool_result(call_id, serialize_tool_result(value)))
+ messages.append(
+ create_tool_result(call_id, serialize_tool_result(value))
+ )
except Exception as e:
call_id = tool_call.get("id", "unknown")
logger.warning("Tool call %s failed: %s", call_id, e)
diff --git a/osmosis_ai/rollout/mcp/loader.py b/osmosis_ai/rollout/mcp/loader.py
index 2ac90f36..a72c19db 100644
--- a/osmosis_ai/rollout/mcp/loader.py
+++ b/osmosis_ai/rollout/mcp/loader.py
@@ -93,11 +93,11 @@ def load_mcp_server(mcp_path: str) -> Any:
# Check fastmcp is installed
try:
- from fastmcp import FastMCP # noqa: F401
- except ImportError:
+ from fastmcp import FastMCP
+ except ImportError as e:
raise MCPLoadError(
"fastmcp is not installed. Install it with: pip install osmosis-ai[mcp]"
- )
+ ) from e
# Import main.py as a submodule of a synthetic package:
# - supports relative imports like `from .tools import ...`
diff --git a/osmosis_ai/rollout/network.py b/osmosis_ai/rollout/network.py
index 05184e4d..3782de99 100644
--- a/osmosis_ai/rollout/network.py
+++ b/osmosis_ai/rollout/network.py
@@ -23,7 +23,6 @@
import ipaddress
import logging
import re
-from typing import Optional, Tuple
import requests
@@ -136,7 +135,7 @@ def is_valid_hostname_or_ip(value: str) -> bool:
# Pattern allows: alphanumeric, hyphens, underscores, dots
# Single char hostnames are valid (e.g., "a", "1")
# Labels can't start or end with hyphen
- hostname_pattern = r'^[A-Za-z0-9_]([A-Za-z0-9_-]{0,61}[A-Za-z0-9_])?(\.[A-Za-z0-9_]([A-Za-z0-9_-]{0,61}[A-Za-z0-9_])?)*$|^[A-Za-z0-9]$'
+ hostname_pattern = r"^[A-Za-z0-9_]([A-Za-z0-9_-]{0,61}[A-Za-z0-9_])?(\.[A-Za-z0-9_]([A-Za-z0-9_-]{0,61}[A-Za-z0-9_])?)*$|^[A-Za-z0-9]$"
return bool(re.match(hostname_pattern, host))
@@ -165,7 +164,7 @@ def is_private_ip(ip: str) -> bool:
# ============================================================================
-def _get_aws_public_ip() -> Optional[str]:
+def _get_aws_public_ip() -> str | None:
"""Get public IP from AWS EC2 IMDSv2 metadata service.
AWS IMDSv2 requires a session token before querying metadata.
@@ -205,7 +204,7 @@ def _get_aws_public_ip() -> Optional[str]:
return None
-def _get_gcp_public_ip() -> Optional[str]:
+def _get_gcp_public_ip() -> str | None:
"""Get public IP from GCP Compute Engine metadata service.
GCP requires the 'Metadata-Flavor: Google' header.
@@ -232,7 +231,7 @@ def _get_gcp_public_ip() -> Optional[str]:
return None
-def _get_azure_public_ip() -> Optional[str]:
+def _get_azure_public_ip() -> str | None:
"""Get public IP from Azure Instance Metadata Service (IMDS).
Note: Only works with Basic SKU public IPs. Standard SKU requires
@@ -261,7 +260,7 @@ def _get_azure_public_ip() -> Optional[str]:
return None
-def detect_from_cloud_metadata() -> Optional[str]:
+def detect_from_cloud_metadata() -> str | None:
"""Try cloud metadata services in parallel.
Cloud metadata services are:
@@ -299,7 +298,9 @@ def detect_from_cloud_metadata() -> Optional[str]:
try:
ip = future.result(timeout=0) # Already completed
if ip:
- logger.debug(f"Cloud metadata detection succeeded: {futures[future]}")
+ logger.debug(
+ f"Cloud metadata detection succeeded: {futures[future]}"
+ )
# Return immediately; other threads will complete on their
# own due to short request timeouts
return ip
@@ -317,7 +318,7 @@ def detect_from_cloud_metadata() -> Optional[str]:
# ============================================================================
-def detect_from_external_services() -> Optional[str]:
+def detect_from_external_services() -> str | None:
"""Get public IP from external services (works on any cloud including Lambda Labs).
Queries multiple services in PARALLEL and returns the first successful result.
@@ -342,7 +343,7 @@ def detect_from_external_services() -> Optional[str]:
("ifconfig.me", "https://ifconfig.me/ip"),
]
- def _query_service(name_url: Tuple[str, str]) -> Optional[str]:
+ def _query_service(name_url: tuple[str, str]) -> str | None:
name, url = name_url
try:
resp = requests.get(url, timeout=_EXTERNAL_SERVICE_TIMEOUT_SECONDS)
@@ -368,7 +369,9 @@ def _query_service(name_url: Tuple[str, str]) -> Optional[str]:
ip = future.result(timeout=0)
if ip:
service_name = futures[future]
- logger.info(f"Detected public IP via external service ({service_name}): {ip}")
+ logger.info(
+ f"Detected public IP via external service ({service_name}): {ip}"
+ )
# Return immediately; other threads will complete on their
# own due to request timeouts
return ip
@@ -442,10 +445,10 @@ def detect_public_ip() -> str:
__all__ = [
- "detect_public_ip",
"PublicIPDetectionError",
+ "detect_public_ip",
+ "is_private_ip",
+ "is_valid_hostname_or_ip",
# Validation helpers
"validate_ipv4",
- "is_valid_hostname_or_ip",
- "is_private_ip",
]
diff --git a/osmosis_ai/rollout/registry.py b/osmosis_ai/rollout/registry.py
index aa48c4c1..c2caac54 100644
--- a/osmosis_ai/rollout/registry.py
+++ b/osmosis_ai/rollout/registry.py
@@ -3,7 +3,7 @@
from __future__ import annotations
import threading
-from typing import Dict, List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
from osmosis_ai.rollout.core.exceptions import AgentLoopNotFoundError
@@ -29,10 +29,10 @@ class AgentLoopRegistry:
def __init__(self):
"""Initialize an empty registry."""
- self._loops: Dict[str, "RolloutAgentLoop"] = {}
+ self._loops: dict[str, RolloutAgentLoop] = {}
self._lock = threading.Lock()
- def register(self, loop: "RolloutAgentLoop") -> None:
+ def register(self, loop: RolloutAgentLoop) -> None:
"""Register an agent loop implementation.
This method is thread-safe.
@@ -68,7 +68,7 @@ def unregister(self, name: str) -> bool:
return True
return False
- def get(self, name: str) -> Optional["RolloutAgentLoop"]:
+ def get(self, name: str) -> RolloutAgentLoop | None:
"""Get an agent loop by name.
This method is thread-safe.
@@ -82,7 +82,7 @@ def get(self, name: str) -> Optional["RolloutAgentLoop"]:
with self._lock:
return self._loops.get(name)
- def list_names(self) -> List[str]:
+ def list_names(self) -> list[str]:
"""List all registered agent loop names.
This method is thread-safe.
@@ -106,7 +106,7 @@ def clear(self) -> None:
_REGISTRY = AgentLoopRegistry()
-def register_agent_loop(loop: "RolloutAgentLoop") -> None:
+def register_agent_loop(loop: RolloutAgentLoop) -> None:
"""Register an agent loop with the global registry.
Args:
@@ -137,7 +137,7 @@ def unregister_agent_loop(name: str) -> bool:
return _REGISTRY.unregister(name)
-def get_agent_loop(name: str) -> "RolloutAgentLoop":
+def get_agent_loop(name: str) -> RolloutAgentLoop:
"""Get an agent loop from the global registry.
Args:
@@ -159,7 +159,7 @@ def get_agent_loop(name: str) -> "RolloutAgentLoop":
return loop
-def list_agent_loops() -> List[str]:
+def list_agent_loops() -> list[str]:
"""List all registered agent loop names in the global registry.
Returns:
diff --git a/osmosis_ai/rollout/server/__init__.py b/osmosis_ai/rollout/server/__init__.py
index 416e541e..292f3699 100644
--- a/osmosis_ai/rollout/server/__init__.py
+++ b/osmosis_ai/rollout/server/__init__.py
@@ -17,7 +17,6 @@
from osmosis_ai.rollout.server.api_key import generate_api_key, validate_api_key
from osmosis_ai.rollout.server.app import create_app
-from osmosis_ai.rollout.server.state import AppState
from osmosis_ai.rollout.server.serve import (
DEFAULT_HOST,
DEFAULT_PORT,
@@ -25,15 +24,16 @@
serve_agent_loop,
validate_and_report,
)
+from osmosis_ai.rollout.server.state import AppState
__all__ = [
- "create_app",
- "AppState",
- "serve_agent_loop",
- "validate_and_report",
- "ServeError",
"DEFAULT_HOST",
"DEFAULT_PORT",
+ "AppState",
+ "ServeError",
+ "create_app",
"generate_api_key",
+ "serve_agent_loop",
+ "validate_and_report",
"validate_api_key",
]
diff --git a/osmosis_ai/rollout/server/api_key.py b/osmosis_ai/rollout/server/api_key.py
index 014c086b..2755865f 100644
--- a/osmosis_ai/rollout/server/api_key.py
+++ b/osmosis_ai/rollout/server/api_key.py
@@ -6,9 +6,8 @@
from __future__ import annotations
-import secrets
import logging
-from typing import Optional
+import secrets
logger = logging.getLogger(__name__)
@@ -36,7 +35,7 @@ def generate_api_key() -> str:
return f"{API_KEY_PREFIX}{random_part}"
-def validate_api_key(provided_key: Optional[str], expected_key: str) -> bool:
+def validate_api_key(provided_key: str | None, expected_key: str) -> bool:
"""Validate a provided API key against the expected key.
Uses constant-time comparison to prevent timing attacks.
@@ -54,7 +53,7 @@ def validate_api_key(provided_key: Optional[str], expected_key: str) -> bool:
__all__ = [
+ "API_KEY_PREFIX",
"generate_api_key",
"validate_api_key",
- "API_KEY_PREFIX",
]
diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py
index d2e561d7..1cbf9de1 100644
--- a/osmosis_ai/rollout/server/app.py
+++ b/osmosis_ai/rollout/server/app.py
@@ -17,20 +17,22 @@
import asyncio
import logging
import time
+from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
-from typing import Any, Awaitable, Callable, Dict, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from fastapi import FastAPI, HTTPException, Request
+
from osmosis_ai.auth.credentials import WorkspaceCredentials
from osmosis_ai.rollout._compat import FASTAPI_AVAILABLE
+from osmosis_ai.rollout.client import OsmosisLLMClient
from osmosis_ai.rollout.config.settings import RolloutSettings, get_settings
from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext
from osmosis_ai.rollout.core.schemas import InitResponse, RolloutRequest
from osmosis_ai.rollout.server.api_key import validate_api_key
from osmosis_ai.rollout.server.state import AppState
-from osmosis_ai.rollout.client import OsmosisLLMClient
logger = logging.getLogger(__name__)
@@ -41,7 +43,7 @@
from fastapi import HTTPException, Request
-def _extract_bearer_token(auth_header: str) -> Optional[str]:
+def _extract_bearer_token(auth_header: str) -> str | None:
"""Extract a bearer token from an Authorization header.
Accepts both:
@@ -60,17 +62,17 @@ def _extract_bearer_token(auth_header: str) -> Optional[str]:
def create_app(
agent_loop: RolloutAgentLoop,
- max_concurrent: Optional[int] = None,
- record_ttl_seconds: Optional[float] = None,
- settings: Optional[RolloutSettings] = None,
- credentials: Optional["WorkspaceCredentials"] = None,
- server_host: Optional[str] = None,
- server_port: Optional[int] = None,
- api_key: Optional[str] = None,
- debug_dir: Optional[str] = None,
- on_startup: Optional[Callable[[], Awaitable[None]]] = None,
- on_shutdown: Optional[Callable[[], Awaitable[None]]] = None,
-) -> "FastAPI":
+ max_concurrent: int | None = None,
+ record_ttl_seconds: float | None = None,
+ settings: RolloutSettings | None = None,
+ credentials: WorkspaceCredentials | None = None,
+ server_host: str | None = None,
+ server_port: int | None = None,
+ api_key: str | None = None,
+ debug_dir: str | None = None,
+ on_startup: Callable[[], Awaitable[None]] | None = None,
+ on_shutdown: Callable[[], Awaitable[None]] | None = None,
+) -> FastAPI:
"""Create a FastAPI application for the agent loop.
This factory creates a complete FastAPI application with:
@@ -167,18 +169,24 @@ async def lifespan(app: FastAPI):
# Start platform registration as a background task.
# The task waits briefly for the server to be ready (after yield),
# then registers with Platform. This ensures the health check succeeds.
- registration_task: Optional[asyncio.Task] = None
- if credentials is not None and server_host is not None and server_port is not None:
+ registration_task: asyncio.Task | None = None
+ if (
+ credentials is not None
+ and server_host is not None
+ and server_port is not None
+ ):
from osmosis_ai.rollout.server.registration import (
- register_with_platform,
print_registration_result,
+ register_with_platform,
)
async def do_registration():
import httpx
# Poll health endpoint until server is ready
- poll_interval = state.settings.registration_readiness_poll_interval_seconds
+ poll_interval = (
+ state.settings.registration_readiness_poll_interval_seconds
+ )
timeout = state.settings.registration_readiness_timeout_seconds
health_url = f"http://127.0.0.1:{server_port}/health"
@@ -259,7 +267,7 @@ async def do_registration():
)
@app.get("/health")
- async def health() -> Dict[str, Any]:
+ async def health() -> dict[str, Any]:
"""Health check endpoint.
Returns server status and statistics.
@@ -272,7 +280,7 @@ async def health() -> Dict[str, Any]:
}
@app.get("/platform/health")
- async def platform_health(request: Request) -> Dict[str, Any]:
+ async def platform_health(request: Request) -> dict[str, Any]:
"""Platform health check endpoint (authenticated).
This endpoint is intended for Osmosis Platform to validate:
@@ -314,13 +322,17 @@ async def init_rollout(
# Validate RolloutServer auth if configured:
# TrainGate must send: Authorization: Bearer
if api_key is not None:
- provided = _extract_bearer_token(http_request.headers.get("authorization") or "")
+ provided = _extract_bearer_token(
+ http_request.headers.get("authorization") or ""
+ )
if not validate_api_key(provided, api_key):
logger.warning(
"Invalid API key for rollout request: rollout_id=%s",
rollout_request.rollout_id,
)
- raise HTTPException(status_code=401, detail="Invalid or missing API key")
+ raise HTTPException(
+ status_code=401, detail="Invalid or missing API key"
+ )
# Idempotency key: prefer request.idempotency_key, fallback to rollout_id.
key = rollout_request.idempotency_key or rollout_request.rollout_id
@@ -415,7 +427,9 @@ async def run_rollout() -> None:
task = asyncio.create_task(run_rollout())
state.mark_started(key, task)
- init_response = InitResponse(rollout_id=rollout_request.rollout_id, tools=tools)
+ init_response = InitResponse(
+ rollout_id=rollout_request.rollout_id, tools=tools
+ )
init_future.set_result(init_response)
logger.info(
diff --git a/osmosis_ai/rollout/server/registration.py b/osmosis_ai/rollout/server/registration.py
index 054a1969..2fcaad2f 100644
--- a/osmosis_ai/rollout/server/registration.py
+++ b/osmosis_ai/rollout/server/registration.py
@@ -8,9 +8,9 @@
import logging
from dataclasses import dataclass
-from typing import Any, Dict, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
-from osmosis_ai.rollout.network import detect_public_ip, PublicIPDetectionError
+from osmosis_ai.rollout.network import PublicIPDetectionError, detect_public_ip
if TYPE_CHECKING:
from osmosis_ai.auth.credentials import WorkspaceCredentials
@@ -23,10 +23,10 @@ class RegistrationResult:
"""Result of platform registration."""
success: bool
- server_id: Optional[str] = None
+ server_id: str | None = None
status: str = "unknown"
- error: Optional[str] = None
- server_info: Optional[Dict[str, Any]] = None
+ error: str | None = None
+ server_info: dict[str, Any] | None = None
@property
def is_healthy(self) -> bool:
@@ -72,8 +72,8 @@ def register_with_platform(
host: str,
port: int,
agent_loop_name: str,
- credentials: "WorkspaceCredentials",
- api_key: Optional[str] = None,
+ credentials: WorkspaceCredentials,
+ api_key: str | None = None,
) -> RegistrationResult:
"""Register the rollout server with Osmosis Platform.
@@ -94,9 +94,9 @@ def register_with_platform(
RegistrationResult with status and any error information.
"""
from osmosis_ai.auth.platform_client import (
- platform_request,
- PlatformAPIError,
AuthenticationExpiredError,
+ PlatformAPIError,
+ platform_request,
)
try:
@@ -117,7 +117,7 @@ def register_with_platform(
)
# Build registration data
- registration_data: Dict[str, Any] = {
+ registration_data: dict[str, Any] = {
"host": report_host,
"port": port,
"agent_loop_name": agent_loop_name,
@@ -183,7 +183,7 @@ def print_registration_result(
host: str,
port: int,
agent_loop_name: str,
- api_key: Optional[str] = None, # noqa: ARG001
+ api_key: str | None = None,
) -> None:
"""Print the registration result to console.
@@ -200,7 +200,7 @@ def print_registration_result(
report_host = host # Fallback to original host for display
if result.is_healthy:
- print(f"\n[OK] Registered with Osmosis Platform")
+ print("\n[OK] Registered with Osmosis Platform")
print(f" Agent: {agent_loop_name}")
print(f" Address: {report_host}:{port}")
print(f" Status: {result.status}")
@@ -209,7 +209,7 @@ def print_registration_result(
print(f" Active rollouts: {active}")
elif result.success:
# Registration succeeded but health check failed
- print(f"\n[WARNING] Registered but health check failed")
+ print("\n[WARNING] Registered but health check failed")
print(f" Agent: {agent_loop_name}")
print(f" Address: {report_host}:{port}")
print(f" Status: {result.status}")
@@ -220,7 +220,7 @@ def print_registration_result(
print(" Tip: Use a VM with public IP and ensure the port is open.")
else:
# Registration failed entirely
- print(f"\n[WARNING] Failed to register with Platform")
+ print("\n[WARNING] Failed to register with Platform")
print(f" Error: {result.error}")
print()
print(" The server will continue running without registration.")
@@ -230,6 +230,6 @@ def print_registration_result(
"RegistrationResult",
"get_public_ip",
"get_report_host",
- "register_with_platform",
"print_registration_result",
+ "register_with_platform",
]
diff --git a/osmosis_ai/rollout/server/serve.py b/osmosis_ai/rollout/server/serve.py
index 510647a6..8d9ebe66 100644
--- a/osmosis_ai/rollout/server/serve.py
+++ b/osmosis_ai/rollout/server/serve.py
@@ -18,22 +18,20 @@ class MyAgent(RolloutAgentLoop):
from __future__ import annotations
import logging
-import sys
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
from osmosis_ai.rollout._compat import FASTAPI_AVAILABLE, UVICORN_AVAILABLE
from osmosis_ai.rollout.console import Console
from osmosis_ai.rollout.core.base import RolloutAgentLoop
from osmosis_ai.rollout.server.api_key import generate_api_key
from osmosis_ai.rollout.validator import (
- AgentLoopValidationError,
ValidationResult,
validate_agent_loop,
)
if TYPE_CHECKING:
- from osmosis_ai.rollout.config.settings import RolloutSettings
from osmosis_ai.auth.credentials import WorkspaceCredentials
+ from osmosis_ai.rollout.config.settings import RolloutSettings
logger = logging.getLogger(__name__)
@@ -55,11 +53,11 @@ def serve_agent_loop(
validate: bool = True,
log_level: str = "info",
reload: bool = False,
- settings: Optional["RolloutSettings"] = None,
+ settings: RolloutSettings | None = None,
skip_register: bool = False,
- api_key: Optional[str] = None,
+ api_key: str | None = None,
local_debug: bool = False,
- debug_dir: Optional[str] = None,
+ debug_dir: str | None = None,
) -> None:
"""Start a RolloutServer for the given agent loop.
@@ -139,7 +137,7 @@ def serve_agent_loop(
)
# Check login status if registration is enabled
- credentials: Optional["WorkspaceCredentials"] = None
+ credentials: WorkspaceCredentials | None = None
if not skip_register:
from osmosis_ai.auth.credentials import get_valid_credentials
@@ -171,7 +169,7 @@ def serve_agent_loop(
logger.info("Using provided API key for server authentication")
# Create debug session directory with timestamp if debug_dir is provided
- debug_session_dir: Optional[str] = None
+ debug_session_dir: str | None = None
if debug_dir:
import os
import time
@@ -273,9 +271,9 @@ def validate_and_report(
return result
-def _get_public_ip() -> Optional[str]:
+def _get_public_ip() -> str | None:
"""Get detected public IP address, or None if detection fails."""
- from osmosis_ai.rollout.network import detect_public_ip, PublicIPDetectionError
+ from osmosis_ai.rollout.network import PublicIPDetectionError, detect_public_ip
try:
return detect_public_ip()
@@ -287,7 +285,9 @@ def _log_validation_result(result: ValidationResult, *, verbose: bool = False) -
"""Log validation result to console."""
console = Console()
if result.valid:
- console.print(f"Agent loop '{result.agent_name}' validated successfully.", style="green")
+ console.print(
+ f"Agent loop '{result.agent_name}' validated successfully.", style="green"
+ )
console.print(f" - Tools: {result.tool_count}")
if result.warnings:
console.print(f" - Warnings: {len(result.warnings)}", style="yellow")
@@ -295,7 +295,9 @@ def _log_validation_result(result: ValidationResult, *, verbose: bool = False) -
for warning in result.warnings:
console.print(f" - {warning}", style="yellow")
else:
- console.print_error(f"Agent loop validation failed with {len(result.errors)} error(s):")
+ console.print_error(
+ f"Agent loop validation failed with {len(result.errors)} error(s):"
+ )
for error in result.errors:
console.print_error(f" - {error}")
if result.warnings and verbose:
diff --git a/osmosis_ai/rollout/server/state.py b/osmosis_ai/rollout/server/state.py
index 36efc42e..43bbfb4b 100644
--- a/osmosis_ai/rollout/server/state.py
+++ b/osmosis_ai/rollout/server/state.py
@@ -20,9 +20,9 @@
from __future__ import annotations
import asyncio
+import contextlib
import logging
import time
-from typing import Dict, Optional, Tuple
from osmosis_ai.rollout.config.settings import RolloutServerSettings, get_settings
from osmosis_ai.rollout.core.schemas import InitResponse
@@ -60,10 +60,10 @@ class AppState:
def __init__(
self,
- max_concurrent: Optional[int] = None,
- record_ttl_seconds: Optional[float] = None,
- cleanup_interval_seconds: Optional[float] = None,
- settings: Optional[RolloutServerSettings] = None,
+ max_concurrent: int | None = None,
+ record_ttl_seconds: float | None = None,
+ cleanup_interval_seconds: float | None = None,
+ settings: RolloutServerSettings | None = None,
agent_loop_name: str = "default",
):
"""Initialize application state.
@@ -81,22 +81,24 @@ def __init__(
self.settings = settings
self._max_concurrent = max_concurrent or settings.max_concurrent_rollouts
self.record_ttl = record_ttl_seconds or settings.record_ttl_seconds
- self._cleanup_interval = cleanup_interval_seconds or settings.cleanup_interval_seconds
+ self._cleanup_interval = (
+ cleanup_interval_seconds or settings.cleanup_interval_seconds
+ )
self._agent_loop_name = agent_loop_name
# NOTE: The "key" used throughout is the idempotency key for init requests:
# - Prefer request.idempotency_key when provided
# - Fallback to request.rollout_id when idempotency_key is missing
- self.rollout_tasks: Dict[str, asyncio.Task] = {}
- self.completed_rollouts: Dict[str, float] = {} # key -> completion_time
+ self.rollout_tasks: dict[str, asyncio.Task] = {}
+ self.completed_rollouts: dict[str, float] = {} # key -> completion_time
# Cached init responses for idempotency (duplicate /v1/rollout/init requests)
- self._init_futures: Dict[str, asyncio.Future[InitResponse]] = {}
+ self._init_futures: dict[str, asyncio.Future[InitResponse]] = {}
self.semaphore = asyncio.Semaphore(self._max_concurrent)
- self._cleanup_task: Optional[asyncio.Task] = None
+ self._cleanup_task: asyncio.Task | None = None
def get_or_create_init_future(
self, key: str
- ) -> Tuple[asyncio.Future[InitResponse], bool]:
+ ) -> tuple[asyncio.Future[InitResponse], bool]:
"""Get or create the init future for a given idempotency key.
This is used to provide true idempotency for /v1/rollout/init:
@@ -135,10 +137,8 @@ async def stop_cleanup_task(self) -> None:
"""
if self._cleanup_task:
self._cleanup_task.cancel()
- try:
+ with contextlib.suppress(asyncio.CancelledError):
await self._cleanup_task
- except asyncio.CancelledError:
- pass
self._cleanup_task = None
logger.info("Cleanup task stopped")
@@ -225,8 +225,14 @@ async def cancel_all(self) -> None:
)
cancelled = sum(1 for r in results if isinstance(r, asyncio.CancelledError))
- errors = sum(1 for r in results if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError))
- logger.info("All rollouts cancelled: cancelled=%d, errors=%d", cancelled, errors)
+ errors = sum(
+ 1
+ for r in results
+ if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError)
+ )
+ logger.info(
+ "All rollouts cancelled: cancelled=%d, errors=%d", cancelled, errors
+ )
# Best-effort cleanup to avoid leaving stale tasks around.
self.rollout_tasks.clear()
diff --git a/osmosis_ai/rollout/testing.py b/osmosis_ai/rollout/testing.py
index 1a3bf817..1bdc3829 100644
--- a/osmosis_ai/rollout/testing.py
+++ b/osmosis_ai/rollout/testing.py
@@ -24,8 +24,9 @@
import threading
import time
+from collections.abc import Callable
from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
import pytest
@@ -33,14 +34,14 @@
from fastapi.testclient import TestClient
from osmosis_ai.rollout.core.schemas import (
- CompletionUsage,
CompletionsRequest,
CompletionsResponse,
+ CompletionUsage,
RolloutResponse,
)
-def fake_token_ids(text: str) -> List[int]:
+def fake_token_ids(text: str) -> list[int]:
"""Generate deterministic fake token IDs for testing.
Args:
@@ -52,7 +53,7 @@ def fake_token_ids(text: str) -> List[int]:
return list(range(len(text)))
-def fake_prompt_token_ids(messages: List[Dict[str, Any]]) -> List[int]:
+def fake_prompt_token_ids(messages: list[dict[str, Any]]) -> list[int]:
"""Generate deterministic fake prompt token IDs for testing.
Token count grows with message count to simulate real behavior.
@@ -87,9 +88,9 @@ class RolloutCompletionTracker:
"""
event: threading.Event = field(default_factory=threading.Event)
- responses: List[Dict[str, Any]] = field(default_factory=list)
+ responses: list[dict[str, Any]] = field(default_factory=list)
- def record(self, response: Dict[str, Any]) -> None:
+ def record(self, response: dict[str, Any]) -> None:
"""Record a completion response and signal the event.
Args:
@@ -115,7 +116,7 @@ def wait(self, timeout: float = 5.0) -> bool:
return self.event.wait(timeout=timeout)
-def _should_use_tools(last_message: Dict[str, Any]) -> bool:
+def _should_use_tools(last_message: dict[str, Any]) -> bool:
"""Determine if the mock trainer should return tool calls.
This heuristic detects calculator-related keywords to trigger tool use.
@@ -136,9 +137,10 @@ def _should_use_tools(last_message: Dict[str, Any]) -> bool:
def create_mock_trainer_app(
- tracker: Optional[RolloutCompletionTracker] = None,
- tool_call_generator: Optional[Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]]] = None,
-) -> "FastAPI":
+ tracker: RolloutCompletionTracker | None = None,
+ tool_call_generator: Callable[[dict[str, Any]], list[dict[str, Any]] | None]
+ | None = None,
+) -> FastAPI:
"""Create a mock trainer FastAPI application for testing.
The mock trainer implements:
@@ -173,16 +175,16 @@ def create_mock_trainer_app(
"""
try:
from fastapi import FastAPI
- except ImportError:
+ except ImportError as e:
raise ImportError(
"FastAPI is required for create_mock_trainer_app(). "
"Install it with: pip install fastapi"
- )
+ ) from e
app = FastAPI(title="Mock Trainer Server")
# In-memory storage of completed rollouts
- completed_rollouts: Dict[str, Dict[str, Any]] = {}
+ completed_rollouts: dict[str, dict[str, Any]] = {}
@app.post("/v1/chat/completions")
async def completions(request: CompletionsRequest) -> CompletionsResponse:
@@ -190,7 +192,7 @@ async def completions(request: CompletionsRequest) -> CompletionsResponse:
last_message = messages[-1] if messages else {"role": "user", "content": ""}
# Determine response based on message content
- tool_calls: Optional[List[Dict[str, Any]]] = None
+ tool_calls: list[dict[str, Any]] | None = None
if tool_call_generator is not None:
tool_calls = tool_call_generator(last_message)
@@ -204,7 +206,7 @@ async def completions(request: CompletionsRequest) -> CompletionsResponse:
]
if tool_calls:
- assistant_message: Dict[str, Any] = {
+ assistant_message: dict[str, Any] = {
"role": "assistant",
"content": "I'll help you with that calculation.",
"tool_calls": tool_calls,
@@ -246,7 +248,7 @@ async def completions(request: CompletionsRequest) -> CompletionsResponse:
)
@app.post("/v1/rollout/completed")
- async def rollout_completed(response: RolloutResponse) -> Dict[str, Any]:
+ async def rollout_completed(response: RolloutResponse) -> dict[str, Any]:
payload = response.model_dump(mode="json", exclude_none=True)
completed_rollouts[response.rollout_id] = payload
@@ -256,19 +258,19 @@ async def rollout_completed(response: RolloutResponse) -> Dict[str, Any]:
return {"status": "ok"}
@app.get("/v1/rollout/completed/{rollout_id}")
- async def get_completed_rollout(rollout_id: str) -> Dict[str, Any]:
+ async def get_completed_rollout(rollout_id: str) -> dict[str, Any]:
return completed_rollouts.get(rollout_id, {})
@app.get("/health")
- async def health() -> Dict[str, Any]:
+ async def health() -> dict[str, Any]:
return {"status": "healthy", "service": "mock-trainer"}
return app
def patch_httpx_for_mock_trainer(
- client: "TestClient",
- monkeypatch: "pytest.MonkeyPatch",
+ client: TestClient,
+ monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Patch httpx.AsyncClient to route requests to the mock trainer.
@@ -313,9 +315,9 @@ async def mock_post(self, url: str, **kwargs):
__all__ = [
- "create_mock_trainer_app",
"RolloutCompletionTracker",
- "patch_httpx_for_mock_trainer",
- "fake_token_ids",
+ "create_mock_trainer_app",
"fake_prompt_token_ids",
+ "fake_token_ids",
+ "patch_httpx_for_mock_trainer",
]
diff --git a/osmosis_ai/rollout/tools.py b/osmosis_ai/rollout/tools.py
index 1fa677fc..c44b5348 100644
--- a/osmosis_ai/rollout/tools.py
+++ b/osmosis_ai/rollout/tools.py
@@ -32,14 +32,15 @@ async def my_executor(tool_call):
import asyncio
import json
import logging
-from typing import Any, Awaitable, Callable, Dict, List, Union
+from collections.abc import Awaitable, Callable
+from typing import Any
-from osmosis_ai.rollout.core.exceptions import ToolArgumentError, ToolExecutionError
+from osmosis_ai.rollout.core.exceptions import ToolArgumentError
logger = logging.getLogger(__name__)
-def create_tool_result(tool_call_id: str, content: str) -> Dict[str, str]:
+def create_tool_result(tool_call_id: str, content: str) -> dict[str, str]:
"""Create a standardized tool result message.
Creates a message dict with role="tool" that can be appended to the
@@ -95,7 +96,7 @@ def serialize_tool_result(result: Any) -> str:
return json.dumps(result)
-def parse_tool_arguments(arguments: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
+def parse_tool_arguments(arguments: str | dict[str, Any]) -> dict[str, Any]:
"""Parse tool call arguments, handling both string and dict formats.
LLM responses may provide arguments as either a JSON string or a dict.
@@ -127,14 +128,14 @@ def parse_tool_arguments(arguments: Union[str, Dict[str, Any]]) -> Dict[str, Any
f"Expected dict from JSON, got {type(parsed).__name__}"
)
except json.JSONDecodeError as e:
- raise ToolArgumentError(f"Invalid JSON arguments: {e}")
+ raise ToolArgumentError(f"Invalid JSON arguments: {e}") from e
raise ToolArgumentError(
f"Expected str or dict arguments, got {type(arguments).__name__}"
)
-def get_tool_call_info(tool_call: Dict[str, Any]) -> tuple[str, str, Dict[str, Any]]:
+def get_tool_call_info(tool_call: dict[str, Any]) -> tuple[str, str, dict[str, Any]]:
"""Extract tool call ID, name, and arguments from a tool call dict.
Args:
@@ -167,15 +168,15 @@ def get_tool_call_info(tool_call: Dict[str, Any]) -> tuple[str, str, Dict[str, A
str(e),
tool_call_id=tool_call_id,
tool_name=function_name,
- )
+ ) from e
return tool_call_id, function_name, arguments
async def execute_tool_calls(
- tool_calls: List[Dict[str, Any]],
- executor: Callable[[Dict[str, Any]], Awaitable[Dict[str, str]]],
-) -> List[Dict[str, str]]:
+ tool_calls: list[dict[str, Any]],
+ executor: Callable[[dict[str, Any]], Awaitable[dict[str, str]]],
+) -> list[dict[str, str]]:
"""Execute multiple tool calls concurrently.
Runs all tool calls in parallel using asyncio.gather. The executor
@@ -207,7 +208,7 @@ async def my_executor(tool_call: Dict) -> Dict[str, str]:
def create_tool_error_result(
tool_call_id: str,
error_message: str,
-) -> Dict[str, str]:
+) -> dict[str, str]:
"""Create an error result message for a failed tool call.
Use this to return error information to the LLM so it can potentially
@@ -230,10 +231,10 @@ def create_tool_error_result(
__all__ = [
+ "create_tool_error_result",
"create_tool_result",
- "serialize_tool_result",
- "parse_tool_arguments",
- "get_tool_call_info",
"execute_tool_calls",
- "create_tool_error_result",
+ "get_tool_call_info",
+ "parse_tool_arguments",
+ "serialize_tool_result",
]
diff --git a/osmosis_ai/rollout/utils.py b/osmosis_ai/rollout/utils.py
index 51282726..ba4d439d 100644
--- a/osmosis_ai/rollout/utils.py
+++ b/osmosis_ai/rollout/utils.py
@@ -15,10 +15,10 @@
from __future__ import annotations
-from typing import Any, Dict, List, Optional
+from typing import Any
-def parse_tool_calls(assistant_message: Dict[str, Any]) -> List[Dict[str, Any]]:
+def parse_tool_calls(assistant_message: dict[str, Any]) -> list[dict[str, Any]]:
"""Safely extract tool_calls from an assistant message.
Handles edge cases where tool_calls may be None, missing, or not a list.
@@ -41,7 +41,7 @@ def parse_tool_calls(assistant_message: Dict[str, Any]) -> List[Dict[str, Any]]:
return []
-def normalize_stop(stop: Any) -> Optional[List[str]]:
+def normalize_stop(stop: Any) -> list[str] | None:
"""Normalize stop parameter to consistent format.
Handles various input formats and normalizes to List[str] or None.
@@ -66,7 +66,7 @@ def normalize_stop(stop: Any) -> Optional[List[str]]:
return None
-def get_message_content(message: Dict[str, Any]) -> str:
+def get_message_content(message: dict[str, Any]) -> str:
"""Extract text content from a message.
Handles messages where content may be None or missing.
@@ -84,7 +84,7 @@ def get_message_content(message: Dict[str, Any]) -> str:
return content if isinstance(content, str) else ""
-def get_message_role(message: Dict[str, Any]) -> str:
+def get_message_role(message: dict[str, Any]) -> str:
"""Extract role from a message.
Args:
@@ -100,7 +100,7 @@ def get_message_role(message: Dict[str, Any]) -> str:
return role if isinstance(role, str) else "unknown"
-def is_assistant_message(message: Dict[str, Any]) -> bool:
+def is_assistant_message(message: dict[str, Any]) -> bool:
"""Check if a message is from the assistant.
Args:
@@ -112,7 +112,7 @@ def is_assistant_message(message: Dict[str, Any]) -> bool:
return get_message_role(message) == "assistant"
-def is_tool_message(message: Dict[str, Any]) -> bool:
+def is_tool_message(message: dict[str, Any]) -> bool:
"""Check if a message is a tool result.
Args:
@@ -124,7 +124,7 @@ def is_tool_message(message: Dict[str, Any]) -> bool:
return get_message_role(message) == "tool"
-def is_user_message(message: Dict[str, Any]) -> bool:
+def is_user_message(message: dict[str, Any]) -> bool:
"""Check if a message is from the user.
Args:
@@ -136,7 +136,7 @@ def is_user_message(message: Dict[str, Any]) -> bool:
return get_message_role(message) == "user"
-def count_messages_by_role(messages: List[Dict[str, Any]]) -> Dict[str, int]:
+def count_messages_by_role(messages: list[dict[str, Any]]) -> dict[str, int]:
"""Count messages by role.
Args:
@@ -149,7 +149,7 @@ def count_messages_by_role(messages: List[Dict[str, Any]]) -> Dict[str, int]:
counts = count_messages_by_role(messages)
# {"user": 3, "assistant": 2, "tool": 1}
"""
- counts: Dict[str, int] = {}
+ counts: dict[str, int] = {}
for message in messages:
role = get_message_role(message)
counts[role] = counts.get(role, 0) + 1
@@ -157,12 +157,12 @@ def count_messages_by_role(messages: List[Dict[str, Any]]) -> Dict[str, int]:
__all__ = [
- "parse_tool_calls",
- "normalize_stop",
+ "count_messages_by_role",
"get_message_content",
"get_message_role",
"is_assistant_message",
"is_tool_message",
"is_user_message",
- "count_messages_by_role",
+ "normalize_stop",
+ "parse_tool_calls",
]
diff --git a/osmosis_ai/rollout/validator.py b/osmosis_ai/rollout/validator.py
index 4696aa1b..943390a9 100644
--- a/osmosis_ai/rollout/validator.py
+++ b/osmosis_ai/rollout/validator.py
@@ -24,7 +24,7 @@ class MyAgent(RolloutAgentLoop):
import logging
from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
from osmosis_ai.rollout.core.base import RolloutAgentLoop
from osmosis_ai.rollout.core.schemas import (
@@ -51,8 +51,8 @@ class ValidationError:
code: str
message: str
- field: Optional[str] = None
- details: Optional[Dict[str, Any]] = None
+ field: str | None = None
+ details: dict[str, Any] | None = None
def __str__(self) -> str:
if self.field:
@@ -73,9 +73,9 @@ class ValidationResult:
"""
valid: bool
- errors: List[ValidationError] = field(default_factory=list)
- warnings: List[ValidationError] = field(default_factory=list)
- agent_name: Optional[str] = None
+ errors: list[ValidationError] = field(default_factory=list)
+ warnings: list[ValidationError] = field(default_factory=list)
+ agent_name: str | None = None
tool_count: int = 0
def __bool__(self) -> bool:
@@ -104,7 +104,7 @@ class AgentLoopValidationError(Exception):
errors: List of validation errors that caused the failure.
"""
- def __init__(self, message: str, errors: Optional[List[ValidationError]] = None):
+ def __init__(self, message: str, errors: list[ValidationError] | None = None):
super().__init__(message)
self.errors = errors or []
@@ -120,9 +120,9 @@ def _create_mock_request() -> RolloutRequest:
)
-def _validate_name(agent_loop: RolloutAgentLoop) -> List[ValidationError]:
+def _validate_name(agent_loop: RolloutAgentLoop) -> list[ValidationError]:
"""Validate agent loop name attribute."""
- errors: List[ValidationError] = []
+ errors: list[ValidationError] = []
name = getattr(agent_loop, "name", None)
if name is None:
@@ -155,14 +155,14 @@ def _validate_name(agent_loop: RolloutAgentLoop) -> List[ValidationError]:
def _validate_tool_schema(
tool: Any, index: int
-) -> tuple[List[ValidationError], List[ValidationError]]:
+) -> tuple[list[ValidationError], list[ValidationError]]:
"""Validate a single tool schema.
Returns:
Tuple of (errors, warnings).
"""
- errors: List[ValidationError] = []
- warnings: List[ValidationError] = []
+ errors: list[ValidationError] = []
+ warnings: list[ValidationError] = []
field_prefix = f"tools[{index}]"
# Check if it's a valid OpenAIFunctionToolSchema or dict
@@ -266,14 +266,14 @@ def _validate_tool_schema(
def _validate_get_tools(
agent_loop: RolloutAgentLoop, request: RolloutRequest
-) -> tuple[List[ValidationError], List[ValidationError], int]:
+) -> tuple[list[ValidationError], list[ValidationError], int]:
"""Validate get_tools() method.
Returns:
Tuple of (errors, warnings, tool_count).
"""
- errors: List[ValidationError] = []
- warnings: List[ValidationError] = []
+ errors: list[ValidationError] = []
+ warnings: list[ValidationError] = []
tool_count = 0
try:
@@ -284,7 +284,10 @@ def _validate_get_tools(
code="GET_TOOLS_EXCEPTION",
message=f"get_tools() raised an exception: {type(e).__name__}: {e}",
field="get_tools",
- details={"exception_type": type(e).__name__, "exception_message": str(e)},
+ details={
+ "exception_type": type(e).__name__,
+ "exception_message": str(e),
+ },
)
)
return errors, warnings, 0
@@ -320,9 +323,9 @@ def _validate_get_tools(
return errors, warnings, tool_count
-def _validate_run_method(agent_loop: RolloutAgentLoop) -> List[ValidationError]:
+def _validate_run_method(agent_loop: RolloutAgentLoop) -> list[ValidationError]:
"""Validate that run() method exists and is async."""
- errors: List[ValidationError] = []
+ errors: list[ValidationError] = []
run_method = getattr(agent_loop, "run", None)
if run_method is None:
@@ -359,7 +362,7 @@ def _validate_run_method(agent_loop: RolloutAgentLoop) -> List[ValidationError]:
def validate_agent_loop(
agent_loop: RolloutAgentLoop,
*,
- request: Optional[RolloutRequest] = None,
+ request: RolloutRequest | None = None,
) -> ValidationResult:
"""Validate a RolloutAgentLoop implementation.
@@ -390,8 +393,8 @@ def validate_agent_loop(
# Or raise exception if invalid
result.raise_if_invalid()
"""
- errors: List[ValidationError] = []
- warnings: List[ValidationError] = []
+ errors: list[ValidationError] = []
+ warnings: list[ValidationError] = []
# Validate name
name_errors = _validate_name(agent_loop)
@@ -421,8 +424,8 @@ def validate_agent_loop(
__all__ = [
+ "AgentLoopValidationError",
"ValidationError",
"ValidationResult",
- "AgentLoopValidationError",
"validate_agent_loop",
]
diff --git a/osmosis_ai/rubric_eval.py b/osmosis_ai/rubric_eval.py
index a668ef7a..5c493fac 100644
--- a/osmosis_ai/rubric_eval.py
+++ b/osmosis_ai/rubric_eval.py
@@ -13,9 +13,15 @@
import os
import re
import warnings
-from typing import Any, Dict, Optional, Union
+from typing import Any
-from .rubric_types import MissingAPIKeyError, ModelInfo, ModelNotFoundError, ProviderRequestError, RewardRubricRunResult
+from .rubric_types import (
+ MissingAPIKeyError,
+ ModelInfo,
+ ModelNotFoundError,
+ ProviderRequestError,
+ RewardRubricRunResult,
+)
# Default timeout for LLM requests
DEFAULT_REQUEST_TIMEOUT_SECONDS = 30.0
@@ -90,7 +96,7 @@ def _default_timeout_for_model(provider: str, model: str) -> float:
# ============================================================================
-def _reward_json_schema() -> Dict[str, Any]:
+def _reward_json_schema() -> dict[str, Any]:
"""Return the JSON schema for rubric evaluation responses."""
return {
"name": "reward_rubric_response",
@@ -136,7 +142,9 @@ def _sanitize_json(raw: str) -> tuple[float, str]:
raise ValueError("Model response must include a finite numeric 'score'.")
if not isinstance(explanation_raw, str) or not explanation_raw.strip():
- raise ValueError("Model response must include a non-empty 'explanation' string.")
+ raise ValueError(
+ "Model response must include a non-empty 'explanation' string."
+ )
return score, explanation_raw.strip()
@@ -158,14 +166,16 @@ def _end_sentinel(label: str) -> str:
return f"<<>>"
-def _quoted_block(label: str, text: Optional[str]) -> str:
+def _quoted_block(label: str, text: str | None) -> str:
if not text or not text.strip():
return ""
cleaned = _escape_triple_backticks(text.strip())
return "\n".join((_start_sentinel(label), cleaned, _end_sentinel(label)))
-def _build_system_prompt(score_min: float, score_max: float, custom_system_prompt: Optional[str]) -> str:
+def _build_system_prompt(
+ score_min: float, score_max: float, custom_system_prompt: str | None
+) -> str:
base = (
"You are an impartial reward judge. "
"Score outputs strictly according to the provided rubric. "
@@ -183,7 +193,7 @@ def _build_system_prompt(score_min: float, score_max: float, custom_system_promp
return base
-def _format_metadata(metadata: Optional[Dict[str, Any]]) -> Optional[str]:
+def _format_metadata(metadata: dict[str, Any] | None) -> str | None:
if not metadata:
return None
try:
@@ -193,7 +203,7 @@ def _format_metadata(metadata: Optional[Dict[str, Any]]) -> Optional[str]:
return json.dumps(serialisable, ensure_ascii=False, indent=2, sort_keys=True)
-def _select_text(*candidates: Optional[str]) -> Optional[str]:
+def _select_text(*candidates: str | None) -> str | None:
for candidate in candidates:
if isinstance(candidate, str):
stripped = candidate.strip()
@@ -207,9 +217,9 @@ def _build_user_prompt(
score_min: float,
score_max: float,
candidate_output: str,
- original_input: Optional[str],
- ground_truth: Optional[str],
- metadata: Optional[Dict[str, Any]],
+ original_input: str | None,
+ ground_truth: str | None,
+ metadata: dict[str, Any] | None,
) -> str:
lines = [
"Rubric:",
@@ -269,7 +279,7 @@ def _build_user_prompt(
# ============================================================================
-def _get_api_key_env_name(provider: str, model_info: ModelInfo) -> Optional[str]:
+def _get_api_key_env_name(provider: str, model_info: ModelInfo) -> str | None:
env_name = model_info.get("api_key_env")
if isinstance(env_name, str):
env_name = env_name.strip()
@@ -278,8 +288,8 @@ def _get_api_key_env_name(provider: str, model_info: ModelInfo) -> Optional[str]
return DEFAULT_API_KEY_ENV.get(provider.lower())
-def _format_api_key_hint(provider: str, env_name: Optional[str]) -> str:
- export_line: Optional[str] = None
+def _format_api_key_hint(provider: str, env_name: str | None) -> str:
+ export_line: str | None = None
if env_name:
export_line = f' export {env_name}="..."'
@@ -291,7 +301,9 @@ def _format_api_key_hint(provider: str, env_name: Optional[str]) -> str:
if export_line:
return "Set the required API key before running:\n\n" + export_line
- exports = "\n".join(f' export {name}="..."' for name in DEFAULT_API_KEY_ENV.values())
+ exports = "\n".join(
+ f' export {name}="..."' for name in DEFAULT_API_KEY_ENV.values()
+ )
return "Set the required API key before running:\n\n" + exports
@@ -351,17 +363,17 @@ def _call_litellm(
score_min: float,
score_max: float,
timeout: float,
- reasoning_effort: Optional[str] = None,
+ reasoning_effort: str | None = None,
) -> RewardRubricRunResult:
"""Call LiteLLM and return the rubric evaluation result."""
try:
import litellm
- except ImportError:
+ except ImportError as e:
raise ProviderRequestError(
provider,
model,
"LiteLLM is required. Install it via `pip install litellm`.",
- )
+ ) from e
# Suppress LiteLLM's "Provider List: ..." debug prints
litellm.suppress_debug_info = True
@@ -377,11 +389,11 @@ def _call_litellm(
# Use json_schema when the model supports it, otherwise fall back to
# json_object (e.g. Cerebras models that only accept the simpler mode).
if litellm.supports_response_schema(model=litellm_model, custom_llm_provider=None):
- response_format: Dict[str, Any] = {"type": "json_schema", "json_schema": schema}
+ response_format: dict[str, Any] = {"type": "json_schema", "json_schema": schema}
else:
response_format = {"type": "json_object"}
- completion_kwargs: Dict[str, Any] = {
+ completion_kwargs: dict[str, Any] = {
"model": litellm_model,
"messages": messages,
"response_format": response_format,
@@ -419,8 +431,16 @@ def _call_litellm(
f"Model '{model}' was not found. Confirm the model identifier is correct "
f"and your {provider} account has access to it.",
) from err
- except (litellm.APIError, litellm.RateLimitError, litellm.AuthenticationError, litellm.Timeout, litellm.APIConnectionError) as err:
- raise ProviderRequestError(provider, model, _extract_error_message(err)) from err
+ except (
+ litellm.APIError,
+ litellm.RateLimitError,
+ litellm.AuthenticationError,
+ litellm.Timeout,
+ litellm.APIConnectionError,
+ ) as err:
+ raise ProviderRequestError(
+ provider, model, _extract_error_message(err)
+ ) from err
except Exception:
# Re-raise other unexpected exceptions to be handled by the caller.
raise
@@ -429,7 +449,9 @@ def _call_litellm(
content = _extract_content(raw)
if not content:
- raise ProviderRequestError(provider, model, "Model response did not include any content.")
+ raise ProviderRequestError(
+ provider, model, "Model response did not include any content."
+ )
try:
score, explanation = _sanitize_json(content)
@@ -457,7 +479,10 @@ def _extract_error_message(err: Exception) -> str:
elif isinstance(error_field, str) and error_field.strip():
return error_field.strip()
- return str(err).strip() or f"{err.__class__.__name__} encountered while contacting provider."
+ return (
+ str(err).strip()
+ or f"{err.__class__.__name__} encountered while contacting provider."
+ )
def _dump_response(response: Any) -> Any:
@@ -472,7 +497,7 @@ def _dump_response(response: Any) -> Any:
return response
-def _extract_content(raw: Any) -> Optional[str]:
+def _extract_content(raw: Any) -> str | None:
"""Extract text content from LiteLLM response."""
if not isinstance(raw, dict):
return None
@@ -500,14 +525,14 @@ def evaluate_rubric(
solution_str: str,
model_info: ModelInfo,
*,
- ground_truth: Optional[str] = None,
- original_input: Optional[str] = None,
- metadata: Optional[Dict[str, Any]] = None,
- score_min: Optional[float] = None,
- score_max: Optional[float] = None,
- timeout: Optional[float] = None,
+ ground_truth: str | None = None,
+ original_input: str | None = None,
+ metadata: dict[str, Any] | None = None,
+ score_min: float | None = None,
+ score_max: float | None = None,
+ timeout: float | None = None,
return_details: bool = False,
-) -> Union[float, RewardRubricRunResult]:
+) -> float | RewardRubricRunResult:
"""
Evaluate a single model output against a rubric by delegating scoring to a hosted LLM.
@@ -547,21 +572,33 @@ def evaluate_rubric(
if not isinstance(solution_str, str) or not solution_str.strip():
raise TypeError("'solution_str' must be a non-empty string")
- resolved_score_min = float(score_min if score_min is not None else model_info.get("score_min", 0.0))
- resolved_score_max = float(score_max if score_max is not None else model_info.get("score_max", 1.0))
+ resolved_score_min = float(
+ score_min if score_min is not None else model_info.get("score_min", 0.0)
+ )
+ resolved_score_max = float(
+ score_max if score_max is not None else model_info.get("score_max", 1.0)
+ )
if resolved_score_max <= resolved_score_min:
raise ValueError("'score_max' must be greater than 'score_min'")
resolved_system_prompt = _select_text(model_info.get("system_prompt"))
- resolved_original_input = _select_text(original_input, model_info.get("original_input"))
+ resolved_original_input = _select_text(
+ original_input, model_info.get("original_input")
+ )
if timeout is not None:
provider_timeout = float(timeout)
else:
model_timeout = model_info.get("timeout")
- provider_timeout = float(model_timeout) if model_timeout else _default_timeout_for_model(provider_name, model)
+ provider_timeout = (
+ float(model_timeout)
+ if model_timeout
+ else _default_timeout_for_model(provider_name, model)
+ )
- system_content = _build_system_prompt(resolved_score_min, resolved_score_max, resolved_system_prompt)
+ system_content = _build_system_prompt(
+ resolved_score_min, resolved_score_max, resolved_system_prompt
+ )
user_content = _build_user_prompt(
rubric,
resolved_score_min,
@@ -589,17 +626,20 @@ def evaluate_rubric(
except ProviderRequestError:
raise
except Exception as exc:
- detail = str(exc).strip() or f"{exc.__class__.__name__} encountered while contacting provider."
+ detail = (
+ str(exc).strip()
+ or f"{exc.__class__.__name__} encountered while contacting provider."
+ )
raise ProviderRequestError(provider_name, model, detail) from exc
return result if return_details else result["score"]
__all__ = [
- "evaluate_rubric",
- "ensure_api_key_available",
+ "DEFAULT_REQUEST_TIMEOUT_SECONDS",
+ "MissingAPIKeyError",
"ModelInfo",
"RewardRubricRunResult",
- "MissingAPIKeyError",
- "DEFAULT_REQUEST_TIMEOUT_SECONDS",
+ "ensure_api_key_available",
+ "evaluate_rubric",
]
diff --git a/osmosis_ai/rubric_types.py b/osmosis_ai/rubric_types.py
index cab2c037..18cc4464 100644
--- a/osmosis_ai/rubric_types.py
+++ b/osmosis_ai/rubric_types.py
@@ -1,6 +1,6 @@
from __future__ import annotations
-from typing import Any, Optional, TypedDict
+from typing import Any, TypedDict
class ModelInfo(TypedDict, total=False):
@@ -10,10 +10,10 @@ class ModelInfo(TypedDict, total=False):
api_key_env: str
score_min: float
score_max: float
- system_prompt: Optional[str]
- original_input: Optional[str]
+ system_prompt: str | None
+ original_input: str | None
timeout: float
- reasoning_effort: Optional[str]
+ reasoning_effort: str | None
class RewardRubricRunResult(TypedDict):
@@ -32,8 +32,14 @@ class ProviderRequestError(RuntimeError):
def __init__(self, provider: str, model: str, detail: str) -> None:
self.provider = provider
self.model = model
- self.detail = detail.strip() if detail else "Provider request failed with no additional detail."
- message = f"Provider '{provider}' request for model '{model}' failed. {self.detail}"
+ self.detail = (
+ detail.strip()
+ if detail
+ else "Provider request failed with no additional detail."
+ )
+ message = (
+ f"Provider '{provider}' request for model '{model}' failed. {self.detail}"
+ )
super().__init__(message)
@@ -42,9 +48,9 @@ class ModelNotFoundError(ProviderRequestError):
__all__ = [
- "ModelInfo",
- "RewardRubricRunResult",
"MissingAPIKeyError",
- "ProviderRequestError",
+ "ModelInfo",
"ModelNotFoundError",
+ "ProviderRequestError",
+ "RewardRubricRunResult",
]
diff --git a/osmosis_ai/utils.py b/osmosis_ai/utils.py
index 7cfc6272..f1295cda 100644
--- a/osmosis_ai/utils.py
+++ b/osmosis_ai/utils.py
@@ -1,9 +1,9 @@
-
import asyncio
import functools
import inspect
import types
-from typing import Any, Callable, Union, get_args, get_origin, get_type_hints
+from collections.abc import Callable
+from typing import Any, Union, get_args, get_origin, get_type_hints
_UNION_TYPE = getattr(types, "UnionType", None)
_ALLOWED_UNION_ORIGINS = (Union,) + ((_UNION_TYPE,) if _UNION_TYPE is not None else ())
@@ -32,49 +32,66 @@ def osmosis_reward(func: Callable) -> Callable:
params = list(sig.parameters.values())
# Classic mode is identified by the first parameter being named 'solution_str'.
- is_classic = len(params) >= 1 and params[0].name == 'solution_str'
+ is_classic = len(params) >= 1 and params[0].name == "solution_str"
if is_classic:
if len(params) < 3:
- raise TypeError(f"Function {func.__name__} must have at least 3 parameters, got {len(params)}")
+ raise TypeError(
+ f"Function {func.__name__} must have at least 3 parameters, got {len(params)}"
+ )
# Check first parameter: solution_str: str
- if params[0].annotation != str:
- raise TypeError(f"First parameter 'solution_str' must be annotated as str, got {params[0].annotation}")
+ if params[0].annotation is not str:
+ raise TypeError(
+ f"First parameter 'solution_str' must be annotated as str, got {params[0].annotation}"
+ )
# Check second parameter: ground_truth: str
- if params[1].name != 'ground_truth':
- raise TypeError(f"Second parameter must be named 'ground_truth', got '{params[1].name}'")
- if params[1].annotation != str:
- raise TypeError(f"Second parameter 'ground_truth' must be annotated as str, got {params[1].annotation}")
+ if params[1].name != "ground_truth":
+ raise TypeError(
+ f"Second parameter must be named 'ground_truth', got '{params[1].name}'"
+ )
+ if params[1].annotation is not str:
+ raise TypeError(
+ f"Second parameter 'ground_truth' must be annotated as str, got {params[1].annotation}"
+ )
# Check third parameter if present: extra_info: dict = None
if len(params) >= 3:
- if params[2].name != 'extra_info':
- raise TypeError(f"Third parameter must be named 'extra_info', got '{params[2].name}'")
- if params[2].annotation != dict:
- raise TypeError(f"Third parameter 'extra_info' must be annotated as dict, got {params[2].annotation}")
+ if params[2].name != "extra_info":
+ raise TypeError(
+ f"Third parameter must be named 'extra_info', got '{params[2].name}'"
+ )
+ if params[2].annotation is not dict:
+ raise TypeError(
+ f"Third parameter 'extra_info' must be annotated as dict, got {params[2].annotation}"
+ )
if params[2].default is inspect.Parameter.empty:
- raise TypeError("Third parameter 'extra_info' must have a default value of None")
+ raise TypeError(
+ "Third parameter 'extra_info' must have a default value of None"
+ )
# Only strip data_source when the function can't accept it.
- _drop_data_source = (
- "data_source" not in sig.parameters
- and not any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
+ _drop_data_source = "data_source" not in sig.parameters and not any(
+ p.kind == inspect.Parameter.VAR_KEYWORD for p in params
)
def _check_classic_return(result):
if is_classic and not isinstance(result, float):
- raise TypeError(f"Function {func.__name__} must return a float, got {type(result).__name__}")
+ raise TypeError(
+ f"Function {func.__name__} must return a float, got {type(result).__name__}"
+ )
return result
if asyncio.iscoroutinefunction(func):
+
@functools.wraps(func)
async def wrapper(*args, **kwargs):
if _drop_data_source:
kwargs.pop("data_source", None)
return _check_classic_return(await func(*args, **kwargs))
else:
+
@functools.wraps(func)
def wrapper(*args, **kwargs):
if _drop_data_source:
@@ -83,6 +100,7 @@ def wrapper(*args, **kwargs):
return wrapper
+
def _is_str_annotation(annotation: Any) -> bool:
if annotation is inspect.Parameter.empty:
return False
@@ -146,41 +164,60 @@ def evaluate_response(
resolved_annotations = {}
if len(params) < 3:
- raise TypeError(f"Function {func.__name__} must have at least 3 parameters, got {len(params)}")
+ raise TypeError(
+ f"Function {func.__name__} must have at least 3 parameters, got {len(params)}"
+ )
solution_param = params[0]
- solution_annotation = resolved_annotations.get(solution_param.name, solution_param.annotation)
+ solution_annotation = resolved_annotations.get(
+ solution_param.name, solution_param.annotation
+ )
if not _is_str_annotation(solution_annotation):
- raise TypeError(f"First parameter 'solution_str' must be annotated as str, got {solution_annotation}")
+ raise TypeError(
+ f"First parameter 'solution_str' must be annotated as str, got {solution_annotation}"
+ )
if solution_param.default is not inspect.Parameter.empty:
- raise TypeError("First parameter 'solution_str' cannot have a default value")
+ raise TypeError(
+ "First parameter 'solution_str' cannot have a default value"
+ )
ground_truth_param = params[1]
if ground_truth_param.name != "ground_truth":
- raise TypeError(f"Second parameter must be named 'ground_truth', got '{ground_truth_param.name}'")
- ground_truth_annotation = resolved_annotations.get(ground_truth_param.name, ground_truth_param.annotation)
+ raise TypeError(
+ f"Second parameter must be named 'ground_truth', got '{ground_truth_param.name}'"
+ )
+ ground_truth_annotation = resolved_annotations.get(
+ ground_truth_param.name, ground_truth_param.annotation
+ )
if not _is_str_annotation(ground_truth_annotation):
union_origin = get_origin(ground_truth_annotation)
if union_origin not in _ALLOWED_UNION_ORIGINS:
raise TypeError(
f"Second parameter 'ground_truth' must be annotated as str or Optional[str], got {ground_truth_annotation}"
)
- union_args = tuple(arg for arg in get_args(ground_truth_annotation) if arg is not type(None)) # noqa: E721
+ union_args = tuple(
+ arg
+ for arg in get_args(ground_truth_annotation)
+ if arg is not type(None)
+ )
if len(union_args) != 1 or not _is_str_annotation(union_args[0]):
raise TypeError(
f"Second parameter 'ground_truth' must be annotated as str or Optional[str], got {ground_truth_annotation}"
)
if ground_truth_param.default is not inspect.Parameter.empty:
- raise TypeError("Second parameter 'ground_truth' cannot have a default value")
+ raise TypeError(
+ "Second parameter 'ground_truth' cannot have a default value"
+ )
extra_info_param = params[2]
if extra_info_param.name != "extra_info":
- raise TypeError(f"Third parameter must be named 'extra_info', got '{extra_info_param.name}'")
+ raise TypeError(
+ f"Third parameter must be named 'extra_info', got '{extra_info_param.name}'"
+ )
# Only strip data_source when the function can't accept it.
- _drop_data_source = (
- "data_source" not in sig.parameters
- and not any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
+ _drop_data_source = "data_source" not in sig.parameters and not any(
+ p.kind == inspect.Parameter.VAR_KEYWORD for p in params
)
def _validate_classic_args(*args, **kwargs):
@@ -191,7 +228,9 @@ def _validate_classic_args(*args, **kwargs):
raise TypeError("'solution_str' argument is required")
solution_value = bound.arguments["solution_str"]
if not isinstance(solution_value, str):
- raise TypeError(f"'solution_str' must be a string, got {type(solution_value).__name__}")
+ raise TypeError(
+ f"'solution_str' must be a string, got {type(solution_value).__name__}"
+ )
if "ground_truth" not in bound.arguments:
raise TypeError("'ground_truth' argument is required")
@@ -206,10 +245,13 @@ def _validate_classic_args(*args, **kwargs):
def _check_classic_return(result):
if is_classic and not isinstance(result, float):
- raise TypeError(f"Function {func.__name__} must return a float, got {type(result).__name__}")
+ raise TypeError(
+ f"Function {func.__name__} must return a float, got {type(result).__name__}"
+ )
return result
if asyncio.iscoroutinefunction(func):
+
@functools.wraps(func)
async def wrapper(*args, **kwargs):
if _drop_data_source:
@@ -218,6 +260,7 @@ async def wrapper(*args, **kwargs):
_validate_classic_args(*args, **kwargs)
return _check_classic_return(await func(*args, **kwargs))
else:
+
@functools.wraps(func)
def wrapper(*args, **kwargs):
if _drop_data_source:
diff --git a/pyproject.toml b/pyproject.toml
index feff8cfa..114c66f1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,8 +48,7 @@ dev = [
"pytest>=8.0.0,<10.0.0",
"pytest-asyncio>=0.23.0,<2.0.0",
"anyio>=4.0.0",
- "black>=25.0.0,<26.0.0",
- "isort>=5.0.0,<7.0.0",
+ "ruff>=0.9.0,<1.0.0",
"fastmcp>=2.0.0",
]
@@ -84,3 +83,28 @@ osmosis_ai = ["py.typed"]
[tool.setuptools.dynamic]
version = {attr = "osmosis_ai.consts.PACKAGE_VERSION"}
+
+[tool.ruff]
+target-version = "py310"
+line-length = 88
+
+[tool.ruff.lint]
+select = [
+ "E", # pycodestyle errors
+ "F", # pyflakes
+ "W", # pycodestyle warnings
+ "I", # isort
+ "UP", # pyupgrade
+ "B", # flake8-bugbear
+ "SIM", # flake8-simplify
+ "RUF", # ruff-specific rules
+]
+ignore = [
+ "E501", # line-too-long (handled by ruff format)
+ "E402", # module-import-not-at-top-of-file (intentional lazy/conditional imports)
+ "SIM117", # multiple-with-statements (nested with is often clearer in tests)
+ "RUF012", # mutable-class-default (common with Pydantic models)
+]
+
+[tool.ruff.lint.isort]
+known-first-party = ["osmosis_ai"]
diff --git a/tests/conftest.py b/tests/conftest.py
index 4422a519..e64219af 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,7 +2,7 @@
from __future__ import annotations
-from typing import Any, Dict, List
+from typing import Any
import pytest
@@ -10,11 +10,11 @@
OpenAIFunctionParametersSchema,
OpenAIFunctionPropertySchema,
OpenAIFunctionSchema,
+ OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
RolloutRequest,
RolloutResult,
- OpenAIFunctionToolSchema,
)
@@ -25,7 +25,7 @@ class MockAgentLoop(RolloutAgentLoop):
def __init__(
self,
- tools: List[OpenAIFunctionToolSchema] | None = None,
+ tools: list[OpenAIFunctionToolSchema] | None = None,
run_result: RolloutResult | None = None,
run_error: Exception | None = None,
):
@@ -33,7 +33,7 @@ def __init__(
self._run_result = run_result
self._run_error = run_error
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return self._tools
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -82,7 +82,7 @@ def sample_tool_schema() -> OpenAIFunctionToolSchema:
@pytest.fixture
-def sample_messages() -> List[Dict[str, Any]]:
+def sample_messages() -> list[dict[str, Any]]:
"""Create sample messages for testing."""
return [
{"role": "system", "content": "You are a helpful assistant."},
@@ -91,13 +91,13 @@ def sample_messages() -> List[Dict[str, Any]]:
@pytest.fixture
-def sample_assistant_message() -> Dict[str, Any]:
+def sample_assistant_message() -> dict[str, Any]:
"""Create a sample assistant message for testing."""
return {"role": "assistant", "content": "The answer is 4."}
@pytest.fixture
-def sample_assistant_message_with_tool_calls() -> Dict[str, Any]:
+def sample_assistant_message_with_tool_calls() -> dict[str, Any]:
"""Create a sample assistant message with tool calls."""
return {
"role": "assistant",
@@ -116,7 +116,7 @@ def sample_assistant_message_with_tool_calls() -> Dict[str, Any]:
@pytest.fixture
-def sample_tool_message() -> Dict[str, Any]:
+def sample_tool_message() -> dict[str, Any]:
"""Create a sample tool result message."""
return {
"role": "tool",
diff --git a/tests/test_auth_credentials.py b/tests/test_auth_credentials.py
index 5f922255..c239cdd1 100644
--- a/tests/test_auth_credentials.py
+++ b/tests/test_auth_credentials.py
@@ -2,7 +2,7 @@
from datetime import datetime, timedelta, timezone
-from osmosis_ai.auth.credentials import WorkspaceCredentials, OrganizationInfo, UserInfo
+from osmosis_ai.auth.credentials import OrganizationInfo, UserInfo, WorkspaceCredentials
def _make_credentials(
@@ -50,5 +50,3 @@ def test_from_dict_rejects_naive_expires_at() -> None:
assert "expires_at must be timezone-aware" in str(exc)
else:
raise AssertionError("Expected ValueError for naive expires_at")
-
-
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 258fbdb6..60406129 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -71,7 +71,12 @@ def test_preview_jsonl(tmp_path, capsys):
"ground_truth": "Assistant cites the subscription policy, clarifies prorated billing, and avoids promising unauthorized discounts.",
"metadata": {"policy_version": "2024-05-01"},
}
- jsonl_content = "\n".join([json.dumps(record_one, ensure_ascii=False), json.dumps(record_two, ensure_ascii=False)])
+ jsonl_content = "\n".join(
+ [
+ json.dumps(record_one, ensure_ascii=False),
+ json.dumps(record_two, ensure_ascii=False),
+ ]
+ )
path = tmp_path / "data.jsonl"
path.write_text(jsonl_content, encoding="utf-8")
@@ -82,7 +87,10 @@ def test_preview_jsonl(tmp_path, capsys):
assert "Loaded 2 JSONL record(s)" in out
assert "JSONL record #1" in out
assert '"rubric_id": "support_followup"' in out
- assert '"ground_truth": "Assistant verifies warranty information, gathers diagnostics, and suggests safe troubleshooting steps."' in out
+ assert (
+ '"ground_truth": "Assistant verifies warranty information, gathers diagnostics, and suggests safe troubleshooting steps."'
+ in out
+ )
def test_eval_command_output_json(tmp_path, monkeypatch, capsys):
@@ -128,8 +136,14 @@ def run(self, config, record):
from osmosis_ai.cli_commands import EvalRubricCommand
fake_evaluator = FakeEvaluator()
- monkeypatch.setattr("osmosis_ai.cli_services.engine.RubricEvaluator", lambda: fake_evaluator)
- monkeypatch.setattr(EvalRubricCommand, "_generate_output_identifier", staticmethod(lambda: "1700000000"))
+ monkeypatch.setattr(
+ "osmosis_ai.cli_services.engine.RubricEvaluator", lambda: fake_evaluator
+ )
+ monkeypatch.setattr(
+ EvalRubricCommand,
+ "_generate_output_identifier",
+ staticmethod(lambda: "1700000000"),
+ )
output_stem = tmp_path / "results"
exit_code = cli.main(
@@ -271,8 +285,14 @@ def run(self, config, record):
from osmosis_ai.cli_commands import EvalRubricCommand
fake_evaluator = FakeEvaluator()
- monkeypatch.setattr("osmosis_ai.cli_services.engine.RubricEvaluator", lambda: fake_evaluator)
- monkeypatch.setattr(EvalRubricCommand, "_generate_output_identifier", staticmethod(lambda: "1700000001"))
+ monkeypatch.setattr(
+ "osmosis_ai.cli_services.engine.RubricEvaluator", lambda: fake_evaluator
+ )
+ monkeypatch.setattr(
+ EvalRubricCommand,
+ "_generate_output_identifier",
+ staticmethod(lambda: "1700000001"),
+ )
output_dir = tmp_path / "baseline_results"
exit_code = cli.main(
@@ -343,7 +363,9 @@ def run(self, config, record):
"raw": {"call": self.calls},
}
- monkeypatch.setattr("osmosis_ai.cli_services.engine.RubricEvaluator", lambda: FakeEvaluator())
+ monkeypatch.setattr(
+ "osmosis_ai.cli_services.engine.RubricEvaluator", lambda: FakeEvaluator()
+ )
output_path = tmp_path / "reports" / "custom_output.txt"
exit_code = cli.main(
@@ -574,7 +596,7 @@ def test_rollout_eval_rejects_batch_size_zero(capsys):
def test_rollout_eval_accepts_any_model_with_base_url(capsys):
"""With --base-url, any model name should be accepted and displayed as-is."""
- exit_code = cli.main(
+ cli.main(
[
"eval",
"-m",
@@ -599,7 +621,7 @@ def test_rollout_eval_accepts_any_model_with_base_url(capsys):
def test_rollout_test_accepts_any_model_with_base_url(capsys):
"""With --base-url, any model name should be accepted and displayed as-is."""
- exit_code = cli.main(
+ cli.main(
[
"test",
"-m",
diff --git a/tests/test_cli_services.py b/tests/test_cli_services.py
index d2f2f397..cdf5a9ca 100644
--- a/tests/test_cli_services.py
+++ b/tests/test_cli_services.py
@@ -46,7 +46,8 @@ def test_baseline_comparator_missing_statistics(tmp_path: Path) -> None:
target.write_text(json.dumps({"metadata": {"average": 0.5}}), encoding="utf-8")
with pytest.raises(
- CLIError, match="Baseline JSON must include an 'overall_statistics' object or top-level statistics."
+ CLIError,
+ match=r"Baseline JSON must include an 'overall_statistics' object or top-level statistics.",
):
comparator.load(target)
@@ -54,9 +55,13 @@ def test_baseline_comparator_missing_statistics(tmp_path: Path) -> None:
def test_baseline_comparator_non_numeric_statistics(tmp_path: Path) -> None:
comparator = BaselineComparator()
target = tmp_path / "baseline.json"
- target.write_text(json.dumps({"overall_statistics": {"average": "bad"}}), encoding="utf-8")
+ target.write_text(
+ json.dumps({"overall_statistics": {"average": "bad"}}), encoding="utf-8"
+ )
- with pytest.raises(CLIError, match="Baseline statistics could not be parsed into numeric values."):
+ with pytest.raises(
+ CLIError, match=r"Baseline statistics could not be parsed into numeric values."
+ ):
comparator.load(target)
@@ -91,11 +96,15 @@ def test_evaluation_session_errors_when_no_matching_records(tmp_path: Path) -> N
session.execute(request)
-def test_resolve_output_path_defaults_to_cache(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+def test_resolve_output_path_defaults_to_cache(
+ monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
monkeypatch.setattr("osmosis_ai.cli_services.session._CACHE_ROOT", tmp_path)
session = EvaluationSession(identifier_factory=lambda: "12345")
- path, identifier = session._resolve_output_path(None, None, rubric_id="My Rubric/ID")
+ path, identifier = session._resolve_output_path(
+ None, None, rubric_id="My Rubric/ID"
+ )
expected_dir = tmp_path / "my_rubric_id"
expected_path = expected_dir / "rubric_eval_result_12345.json"
@@ -106,7 +115,9 @@ def test_resolve_output_path_defaults_to_cache(monkeypatch: pytest.MonkeyPatch,
def _write_records(path: Path, lines: list[dict]) -> None:
- path.write_text("\n".join(json.dumps(line) for line in lines) + "\n", encoding="utf-8")
+ path.write_text(
+ "\n".join(json.dumps(line) for line in lines) + "\n", encoding="utf-8"
+ )
def test_dataset_loader_invalid_json(tmp_path: Path) -> None:
@@ -133,7 +144,10 @@ def test_dataset_loader_missing_solution_str(tmp_path: Path) -> None:
record = {"rubric_id": "support_followup"}
_write_records(data_path, [record])
- with pytest.raises(CLIError, match="Record 'support_followup' must include a non-empty 'solution_str' string."):
+ with pytest.raises(
+ CLIError,
+ match=r"Record 'support_followup' must include a non-empty 'solution_str' string.",
+ ):
loader.load(data_path)
@@ -146,7 +160,10 @@ def test_dataset_loader_blank_solution_str(tmp_path: Path) -> None:
}
_write_records(data_path, [record])
- with pytest.raises(CLIError, match="Record 'support_followup' must include a non-empty 'solution_str' string."):
+ with pytest.raises(
+ CLIError,
+ match=r"Record 'support_followup' must include a non-empty 'solution_str' string.",
+ ):
loader.load(data_path)
@@ -172,7 +189,9 @@ def test_dataset_loader_supports_original_input_in_extra_info(tmp_path: Path) ->
loaded = loader.load(data_path)[0]
assert loaded.original_input == "Please help me troubleshoot my purifier."
- assert loaded.extra_info == {"original_input": "Please help me troubleshoot my purifier."}
+ assert loaded.extra_info == {
+ "original_input": "Please help me troubleshoot my purifier."
+ }
def test_rubric_config_rejects_extra_info(tmp_path: Path) -> None:
@@ -227,7 +246,9 @@ def test_dataset_record_assistant_preview_truncates(tmp_path: Path) -> None:
assert preview == ("A" * 137) + "..."
-def test_dataset_record_assistant_preview_returns_none_without_assistant(tmp_path: Path) -> None:
+def test_dataset_record_assistant_preview_returns_none_without_assistant(
+ tmp_path: Path,
+) -> None:
record = DatasetRecord(
payload={},
rubric_id="support_followup",
@@ -271,7 +292,12 @@ def fake_evaluate(**kwargs):
config = RubricConfig(
rubric_id="support_followup",
rubric_text="Score how well the assistant resolves the issue.",
- model_info={"provider": "openai", "model": "gpt-5-mini", "api_key": "dummy", "system_prompt": "Judge fairly."},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5-mini",
+ "api_key": "dummy",
+ "system_prompt": "Judge fairly.",
+ },
score_min=0.0,
score_max=1.0,
system_prompt="System override prompt.",
diff --git a/tests/test_eval_fn.py b/tests/test_eval_fn.py
index d1664292..b7f77817 100644
--- a/tests/test_eval_fn.py
+++ b/tests/test_eval_fn.py
@@ -11,7 +11,9 @@ def _dummy_eval(solution_str: str, ground_truth: str, extra_info: dict) -> float
return 1.0
-def test_load_eval_fns_uses_module_path_as_name(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_load_eval_fns_uses_module_path_as_name(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
monkeypatch.setattr(
"osmosis_ai.rollout.eval.evaluation.eval_fn.load_eval_fn",
lambda _path: _dummy_eval,
diff --git a/tests/test_eval_report.py b/tests/test_eval_report.py
index 77a18627..ff6445bd 100644
--- a/tests/test_eval_report.py
+++ b/tests/test_eval_report.py
@@ -4,8 +4,8 @@
from io import StringIO
-from osmosis_ai.rollout.console import Console
import osmosis_ai.rollout.eval.evaluation.report as report_module
+from osmosis_ai.rollout.console import Console
from osmosis_ai.rollout.eval.evaluation.runner import EvalEvalSummary, EvalResult
diff --git a/tests/test_eval_runner.py b/tests/test_eval_runner.py
index 0de0951b..559b8beb 100644
--- a/tests/test_eval_runner.py
+++ b/tests/test_eval_runner.py
@@ -3,7 +3,7 @@
from __future__ import annotations
import asyncio
-from typing import Any, Dict, List
+from typing import Any
import pytest
@@ -16,11 +16,11 @@
RolloutRequest,
RolloutResult,
)
-from osmosis_ai.rollout.eval.evaluation.eval_fn import EvalFnWrapper
-from osmosis_ai.rollout.eval.evaluation.runner import EvalRunner
from osmosis_ai.rollout.client import CompletionsResult
from osmosis_ai.rollout.core.schemas import RolloutMetrics
from osmosis_ai.rollout.eval.common.dataset import DatasetRow
+from osmosis_ai.rollout.eval.evaluation.eval_fn import EvalFnWrapper
+from osmosis_ai.rollout.eval.evaluation.runner import EvalRunner
class MockLLMClient:
@@ -29,7 +29,7 @@ def __init__(self, model: str = "mock-model") -> None:
self.display_name = model
self._api_key: str | None = None
self._api_base: str | None = None
- self._tools: List[Dict[str, Any]] | None = None
+ self._tools: list[dict[str, Any]] | None = None
self._prompt_tokens = 0
self._response_tokens = 0
self._num_llm_calls = 0
@@ -41,7 +41,7 @@ def __init__(self, model: str = "mock-model") -> None:
finish_reason="stop",
)
- def set_tools(self, tools: List[Any]) -> None:
+ def set_tools(self, tools: list[Any]) -> None:
if tools:
self._tools = [
t.model_dump(exclude_none=True) if hasattr(t, "model_dump") else t
@@ -67,7 +67,7 @@ def get_metrics(self) -> RolloutMetrics:
)
async def chat_completions(
- self, messages: List[Dict[str, Any]], **kwargs: Any
+ self, messages: list[dict[str, Any]], **kwargs: Any
) -> CompletionsResult:
if self._tools is not None and "tools" not in kwargs:
kwargs["tools"] = self._tools
@@ -82,7 +82,7 @@ class MockAgentLoop(RolloutAgentLoop):
def __init__(
self,
- tools: List[OpenAIFunctionToolSchema] | None = None,
+ tools: list[OpenAIFunctionToolSchema] | None = None,
run_error: Exception | None = None,
call_llm: bool = False,
) -> None:
@@ -90,7 +90,7 @@ def __init__(
self._run_error = run_error
self._call_llm = call_llm
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return self._tools
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -133,9 +133,9 @@ async def test_run_single_applies_eval_functions(self) -> None:
agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
async def full_eval(
- messages: List[Dict[str, Any]],
+ messages: list[dict[str, Any]],
ground_truth: str,
- metadata: Dict[str, Any],
+ metadata: dict[str, Any],
) -> float:
assert ground_truth.startswith("Answer")
assert "user_prompt" in metadata
@@ -144,7 +144,7 @@ async def full_eval(
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0 if "response" in solution_str else 0.0
@@ -179,7 +179,7 @@ async def test_run_single_propagates_agent_failure(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -211,7 +211,7 @@ async def test_run_eval_computes_pass_at_k(self) -> None:
def alternating_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
call_counter["n"] += 1
return 1.0 if call_counter["n"] % 2 == 1 else 0.0
@@ -243,7 +243,7 @@ async def test_run_eval_concurrent(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0 if "response" in solution_str else 0.0
@@ -285,7 +285,7 @@ async def test_run_eval_batch_size_gt_one_defaults_to_concurrent(
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0 if "response" in solution_str else 0.0
@@ -332,7 +332,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -372,7 +372,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -399,6 +399,7 @@ def simple_eval(
@pytest.mark.asyncio
async def test_run_eval_concurrent_normal_failures_continue(self) -> None:
"""Non-systemic concurrent failures should be recorded without early stopping."""
+
class FailingAgent(MockAgentLoop):
async def run(self, ctx: RolloutContext) -> RolloutResult:
raise RuntimeError("eval failure")
@@ -409,7 +410,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -443,7 +444,7 @@ async def test_run_eval_counts_failed_runs_as_zero_scores(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -526,7 +527,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -547,7 +548,9 @@ def simple_eval(
assert first_run.tokens == 15
@pytest.mark.asyncio
- async def test_run_eval_concurrent_systemic_error_preserves_duration_and_tokens(self) -> None:
+ async def test_run_eval_concurrent_systemic_error_preserves_duration_and_tokens(
+ self,
+ ) -> None:
"""Concurrent systemic failures should keep per-run duration/token stats."""
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
@@ -564,7 +567,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -581,7 +584,9 @@ def simple_eval(
batch_size=2,
)
- failed_runs = [run for row in eval_result.rows for run in row.runs if not run.success]
+ failed_runs = [
+ run for row in eval_result.rows for run in row.runs if not run.success
+ ]
assert failed_runs
assert all(run.duration_ms > 0 for run in failed_runs)
assert all(run.tokens == 15 for run in failed_runs)
@@ -598,7 +603,7 @@ async def test_run_eval_continues_on_non_systemic_failure(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -632,7 +637,7 @@ async def test_run_single_with_model_tag(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -668,7 +673,7 @@ async def test_run_eval_baseline_sequential(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -712,7 +717,7 @@ async def test_run_eval_baseline_concurrent(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -752,7 +757,7 @@ async def test_run_eval_no_baseline_backward_compat(self) -> None:
def simple_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0
@@ -789,7 +794,7 @@ async def test_eval_summaries_not_polluted_by_baseline_sequential(self) -> None:
def score_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
# Primary gets "eval response" -> 1.0, baseline gets "baseline wrong answer" -> 0.0
return 1.0 if "eval response" in solution_str else 0.0
@@ -834,7 +839,7 @@ async def test_eval_summaries_not_polluted_by_baseline_concurrent(self) -> None:
def score_eval(
solution_str: str,
ground_truth: str,
- extra_info: Dict[str, Any],
+ extra_info: dict[str, Any],
) -> float:
return 1.0 if "eval response" in solution_str else 0.0
diff --git a/tests/test_litellm_provider.py b/tests/test_litellm_provider.py
index 4e725b8d..4a9d5075 100644
--- a/tests/test_litellm_provider.py
+++ b/tests/test_litellm_provider.py
@@ -8,11 +8,11 @@
import pytest
from osmosis_ai.rubric_eval import (
- _to_litellm_model,
+ DEFAULT_REQUEST_TIMEOUT_SECONDS,
_default_timeout_for_model,
_sanitize_json,
+ _to_litellm_model,
evaluate_rubric,
- DEFAULT_REQUEST_TIMEOUT_SECONDS,
)
from osmosis_ai.rubric_types import ModelNotFoundError, ProviderRequestError
@@ -21,9 +21,13 @@ class TestToLitellmModel:
"""Tests for model format conversion."""
def test_openai_gpt5_uses_responses_prefix(self):
- assert _to_litellm_model("openai", "gpt-5-mini") == "openai/responses/gpt-5-mini"
+ assert (
+ _to_litellm_model("openai", "gpt-5-mini") == "openai/responses/gpt-5-mini"
+ )
assert _to_litellm_model("openai", "gpt-5") == "openai/responses/gpt-5"
- assert _to_litellm_model("OpenAI", "GPT-5-turbo") == "openai/responses/GPT-5-turbo"
+ assert (
+ _to_litellm_model("OpenAI", "GPT-5-turbo") == "openai/responses/GPT-5-turbo"
+ )
def test_openai_other_models_standard_prefix(self):
assert _to_litellm_model("openai", "gpt-4o") == "openai/gpt-4o"
@@ -31,27 +35,44 @@ def test_openai_other_models_standard_prefix(self):
assert _to_litellm_model("openai", "o1-preview") == "openai/o1-preview"
def test_anthropic_models(self):
- assert _to_litellm_model("anthropic", "claude-sonnet-4-5-20250929") == "anthropic/claude-sonnet-4-5-20250929"
- assert _to_litellm_model("anthropic", "claude-3-opus-20240229") == "anthropic/claude-3-opus-20240229"
+ assert (
+ _to_litellm_model("anthropic", "claude-sonnet-4-5-20250929")
+ == "anthropic/claude-sonnet-4-5-20250929"
+ )
+ assert (
+ _to_litellm_model("anthropic", "claude-3-opus-20240229")
+ == "anthropic/claude-3-opus-20240229"
+ )
def test_xai_models(self):
assert _to_litellm_model("xai", "grok-4-fast") == "xai/grok-4-fast"
assert _to_litellm_model("xai", "grok-2") == "xai/grok-2"
def test_gemini_models(self):
- assert _to_litellm_model("gemini", "gemini-2.0-flash") == "gemini/gemini-2.0-flash"
+ assert (
+ _to_litellm_model("gemini", "gemini-2.0-flash") == "gemini/gemini-2.0-flash"
+ )
assert _to_litellm_model("gemini", "gemini-1.5-pro") == "gemini/gemini-1.5-pro"
def test_other_providers(self):
- assert _to_litellm_model("cerebras", "llama-3.1-70b") == "cerebras/llama-3.1-70b"
- assert _to_litellm_model("openrouter", "meta-llama/llama-3-70b") == "openrouter/meta-llama/llama-3-70b"
+ assert (
+ _to_litellm_model("cerebras", "llama-3.1-70b") == "cerebras/llama-3.1-70b"
+ )
+ assert (
+ _to_litellm_model("openrouter", "meta-llama/llama-3-70b")
+ == "openrouter/meta-llama/llama-3-70b"
+ )
def test_empty_provider_returns_model_only(self):
assert _to_litellm_model("", "gpt-4o") == "gpt-4o"
def test_case_insensitive(self):
- assert _to_litellm_model("OPENAI", "gpt-5-mini") == "openai/responses/gpt-5-mini"
- assert _to_litellm_model("Anthropic", "claude-3-opus") == "anthropic/claude-3-opus"
+ assert (
+ _to_litellm_model("OPENAI", "gpt-5-mini") == "openai/responses/gpt-5-mini"
+ )
+ assert (
+ _to_litellm_model("Anthropic", "claude-3-opus") == "anthropic/claude-3-opus"
+ )
class TestDefaultTimeoutForModel:
@@ -69,7 +90,10 @@ def test_openai_gpt5_timeout(self):
assert _default_timeout_for_model("openai", "gpt-5") == 45.0
def test_openai_other_timeout(self):
- assert _default_timeout_for_model("openai", "gpt-4o") == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ assert (
+ _default_timeout_for_model("openai", "gpt-4o")
+ == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ )
def test_gemini_timeout(self):
assert _default_timeout_for_model("gemini", "gemini-2.0-flash") == 45.0
@@ -78,22 +102,32 @@ def test_cerebras_timeout(self):
assert _default_timeout_for_model("cerebras", "llama-4-scout-17b") == 60.0
def test_openrouter_timeout(self):
- assert _default_timeout_for_model("openrouter", "meta-llama/llama-4-maverick") == 60.0
+ assert (
+ _default_timeout_for_model("openrouter", "meta-llama/llama-4-maverick")
+ == 60.0
+ )
def test_anthropic_timeout(self):
- assert _default_timeout_for_model("anthropic", "claude-3-opus") == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ assert (
+ _default_timeout_for_model("anthropic", "claude-3-opus")
+ == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ )
class TestSanitizeJson:
"""Tests for JSON response parsing."""
def test_valid_json(self):
- score, explanation = _sanitize_json('{"score": 0.85, "explanation": "Good response"}')
+ score, explanation = _sanitize_json(
+ '{"score": 0.85, "explanation": "Good response"}'
+ )
assert score == 0.85
assert explanation == "Good response"
def test_json_with_code_fence(self):
- score, explanation = _sanitize_json('```json\n{"score": 0.9, "explanation": "Test"}\n```')
+ score, explanation = _sanitize_json(
+ '```json\n{"score": 0.9, "explanation": "Test"}\n```'
+ )
assert score == 0.9
assert explanation == "Test"
@@ -130,13 +164,19 @@ class TestEvaluateRubric:
def test_evaluate_rubric_success(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(0.85, "Good response")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.85, "Good response"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score accuracy",
solution_str="The answer is 42",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert result == 0.85
@@ -148,13 +188,19 @@ def test_evaluate_rubric_success(self):
def test_evaluate_rubric_gpt5_no_temperature(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(0.9, "Excellent")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.9, "Excellent"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
evaluate_rubric(
rubric="Score quality",
solution_str="Test response",
- model_info={"provider": "openai", "model": "gpt-5-mini", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5-mini",
+ "api_key": "test-key",
+ },
)
call_kwargs = mock_litellm.completion.call_args.kwargs
@@ -163,13 +209,19 @@ def test_evaluate_rubric_gpt5_no_temperature(self):
def test_evaluate_rubric_anthropic(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(0.9, "Excellent")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.9, "Excellent"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score quality",
solution_str="Well written response",
- model_info={"provider": "anthropic", "model": "claude-sonnet-4-5-20250929", "api_key": "test-key"},
+ model_info={
+ "provider": "anthropic",
+ "model": "claude-sonnet-4-5-20250929",
+ "api_key": "test-key",
+ },
)
assert result == 0.9
@@ -178,13 +230,19 @@ def test_evaluate_rubric_anthropic(self):
def test_evaluate_rubric_xai(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(0.8, "Good")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.8, "Good"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score response",
solution_str="Some response",
- model_info={"provider": "xai", "model": "grok-4-fast", "api_key": "test-key"},
+ model_info={
+ "provider": "xai",
+ "model": "grok-4-fast",
+ "api_key": "test-key",
+ },
)
assert result == 0.8
@@ -193,13 +251,19 @@ def test_evaluate_rubric_xai(self):
def test_evaluate_rubric_return_details(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(0.65, "Detailed explanation")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.65, "Detailed explanation"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score accuracy",
solution_str="Some answer",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
return_details=True,
)
@@ -209,26 +273,38 @@ def test_evaluate_rubric_return_details(self):
def test_evaluate_rubric_score_clamping_max(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(1.5, "Over max")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 1.5, "Over max"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert result == 1.0 # Clamped to max
def test_evaluate_rubric_score_clamping_min(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(-0.5, "Under min")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ -0.5, "Under min"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert result == 0.0 # Clamped to min
@@ -253,7 +329,11 @@ def test_evaluate_rubric_api_error(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert "rate limit" in str(exc_info.value).lower()
@@ -271,7 +351,11 @@ def test_evaluate_rubric_invalid_json_response(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert "not valid JSON" in str(exc_info.value)
@@ -289,7 +373,11 @@ def test_evaluate_rubric_empty_response(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert "did not include any content" in str(exc_info.value)
@@ -314,7 +402,11 @@ def test_evaluate_rubric_model_not_found_404(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "no-such-model", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "no-such-model",
+ "api_key": "test-key",
+ },
)
assert exc_info.value.provider == "openai"
@@ -340,7 +432,11 @@ def test_evaluate_rubric_model_not_found_is_provider_request_error(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "anthropic", "model": "bad-model", "api_key": "test-key"},
+ model_info={
+ "provider": "anthropic",
+ "model": "bad-model",
+ "api_key": "test-key",
+ },
)
def test_evaluate_rubric_non_404_error_is_not_model_not_found(self):
@@ -363,20 +459,30 @@ def test_evaluate_rubric_non_404_error_is_not_model_not_found(self):
evaluate_rubric(
rubric="Score",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
assert not isinstance(exc_info.value, ModelNotFoundError)
def test_evaluate_rubric_custom_score_range(self):
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(7.5, "Good")
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 7.5, "Good"
+ )
with patch.dict(sys.modules, {"litellm": mock_litellm}):
result = evaluate_rubric(
rubric="Score from 0-10",
solution_str="Response",
- model_info={"provider": "openai", "model": "gpt-4o", "api_key": "test-key"},
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
score_min=0.0,
score_max=10.0,
)
diff --git a/tests/test_mcp_agent_loop.py b/tests/test_mcp_agent_loop.py
index d2b97b7d..25bc848d 100644
--- a/tests/test_mcp_agent_loop.py
+++ b/tests/test_mcp_agent_loop.py
@@ -3,13 +3,11 @@
from __future__ import annotations
import argparse
-import os
import sys
import textwrap
import types
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -23,13 +21,12 @@
)
from osmosis_ai.rollout.mcp.loader import MCPLoadError, load_mcp_server
-
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
-def _make_mock_mcp_server(tools: Dict[str, Dict[str, Any]]) -> MagicMock:
+def _make_mock_mcp_server(tools: dict[str, dict[str, Any]]) -> MagicMock:
"""Create a mock FastMCP server with given tools.
Args:
@@ -40,7 +37,9 @@ def _make_mock_mcp_server(tools: Dict[str, Dict[str, Any]]) -> MagicMock:
tool = MagicMock()
tool.name = name
tool.description = info.get("description", "")
- tool.parameters = info.get("parameters", {"type": "object", "properties": {}, "required": []})
+ tool.parameters = info.get(
+ "parameters", {"type": "object", "properties": {}, "required": []}
+ )
mock_tools[name] = tool
tool_manager = MagicMock()
@@ -248,12 +247,18 @@ def test_custom_agent_name(self):
class TestMCPAgentLoopRun:
async def test_no_tool_calls(self):
"""LLM responds without tool calls — loop should complete immediately."""
- server = _make_mock_mcp_server({
- "add": {
- "description": "Add",
- "parameters": {"type": "object", "properties": {"a": {"type": "integer"}}, "required": ["a"]},
- },
- })
+ server = _make_mock_mcp_server(
+ {
+ "add": {
+ "description": "Add",
+ "parameters": {
+ "type": "object",
+ "properties": {"a": {"type": "integer"}},
+ "required": ["a"],
+ },
+ },
+ }
+ )
agent = MCPAgentLoop(server)
# Build context
@@ -294,16 +299,21 @@ async def test_no_tool_calls(self):
async def test_tool_call_then_response(self):
"""LLM makes a tool call, then responds with final answer."""
# Set up MCP server with a tool
- server = _make_mock_mcp_server({
- "add": {
- "description": "Add",
- "parameters": {
- "type": "object",
- "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
- "required": ["a", "b"],
+ server = _make_mock_mcp_server(
+ {
+ "add": {
+ "description": "Add",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "a": {"type": "integer"},
+ "b": {"type": "integer"},
+ },
+ "required": ["a", "b"],
+ },
},
- },
- })
+ }
+ )
# Mock call_tool to return a ToolResult-like object
tool_result = MagicMock()
@@ -370,12 +380,14 @@ async def test_tool_call_then_response(self):
async def test_tool_call_exception_handled(self):
"""Tool execution error should be captured as error message, loop continues."""
- server = _make_mock_mcp_server({
- "fail_tool": {
- "description": "Always fails",
- "parameters": {"type": "object", "properties": {}, "required": []},
- },
- })
+ server = _make_mock_mcp_server(
+ {
+ "fail_tool": {
+ "description": "Always fails",
+ "parameters": {"type": "object", "properties": {}, "required": []},
+ },
+ }
+ )
server._tool_manager.call_tool = AsyncMock(side_effect=RuntimeError("boom"))
agent = MCPAgentLoop(server)
@@ -402,7 +414,11 @@ async def test_tool_call_exception_handled(self):
"role": "assistant",
"content": None,
"tool_calls": [
- {"id": "call_err", "type": "function", "function": {"name": "fail_tool", "arguments": "{}"}},
+ {
+ "id": "call_err",
+ "type": "function",
+ "function": {"name": "fail_tool", "arguments": "{}"},
+ },
],
}
tc_resp.has_tool_calls = True
@@ -477,7 +493,7 @@ def test_directory_not_found(self):
load_mcp_server("/nonexistent/path/to/mcp")
def test_no_main_py(self, tmp_path):
- with pytest.raises(MCPLoadError, match="No main.py found"):
+ with pytest.raises(MCPLoadError, match=r"No main\.py found"):
load_mcp_server(str(tmp_path))
def test_no_fastmcp_instance(self, tmp_path):
@@ -488,7 +504,8 @@ def test_no_fastmcp_instance(self, tmp_path):
def test_successful_load(self, tmp_path):
main_py = tmp_path / "main.py"
- main_py.write_text(textwrap.dedent("""\
+ main_py.write_text(
+ textwrap.dedent("""\
from fastmcp import FastMCP
mcp = FastMCP("test_server")
@@ -497,11 +514,13 @@ def test_successful_load(self, tmp_path):
def add(a: int, b: int) -> int:
\"\"\"Add two numbers\"\"\"
return a + b
- """))
+ """)
+ )
server = load_mcp_server(str(tmp_path))
from fastmcp import FastMCP
+
assert isinstance(server, FastMCP)
assert "add" in server._tool_manager._tools
@@ -515,10 +534,12 @@ def test_restores_sys_path_after_load(self, tmp_path):
original_sys_path = list(sys.path)
main_py = tmp_path / "main.py"
- main_py.write_text(textwrap.dedent("""\
+ main_py.write_text(
+ textwrap.dedent("""\
from fastmcp import FastMCP
mcp = FastMCP("path_test")
- """))
+ """)
+ )
_ = load_mcp_server(str(tmp_path))
@@ -526,7 +547,8 @@ def test_restores_sys_path_after_load(self, tmp_path):
def test_supports_relative_imports_in_main(self, tmp_path):
helpers_py = tmp_path / "helpers.py"
- helpers_py.write_text(textwrap.dedent("""\
+ helpers_py.write_text(
+ textwrap.dedent("""\
from fastmcp import FastMCP
def build_server():
@@ -537,14 +559,17 @@ def ping() -> str:
return "pong"
return mcp
- """))
+ """)
+ )
main_py = tmp_path / "main.py"
- main_py.write_text(textwrap.dedent("""\
+ main_py.write_text(
+ textwrap.dedent("""\
from .helpers import build_server
mcp = build_server()
- """))
+ """)
+ )
server = load_mcp_server(str(tmp_path))
assert "ping" in server._tool_manager._tools
@@ -558,20 +583,24 @@ def test_failed_load_does_not_leak_absolute_sibling_modules(self, tmp_path):
first_helper = first_dir / "leaky_shared_mod.py"
first_helper.write_text('VALUE = "from_first"\n')
first_main = first_dir / "main.py"
- first_main.write_text(textwrap.dedent("""\
+ first_main.write_text(
+ textwrap.dedent("""\
import leaky_shared_mod
raise RuntimeError("boom after import")
- """))
+ """)
+ )
second_helper = second_dir / "leaky_shared_mod.py"
second_helper.write_text('VALUE = "from_second"\n')
second_main = second_dir / "main.py"
- second_main.write_text(textwrap.dedent("""\
+ second_main.write_text(
+ textwrap.dedent("""\
import leaky_shared_mod
from fastmcp import FastMCP
mcp = FastMCP(leaky_shared_mod.VALUE)
- """))
+ """)
+ )
with pytest.raises(MCPLoadError, match="Error importing"):
load_mcp_server(str(first_dir))
@@ -708,13 +737,13 @@ async def _run_eval_cli_and_capture_run_eval_kwargs(
self,
monkeypatch: pytest.MonkeyPatch,
batch_size: int,
- ) -> Dict[str, Any]:
+ ) -> dict[str, Any]:
from osmosis_ai.rollout.eval.evaluation.cli import EvalCommand
- captured_kwargs: Dict[str, Any] = {}
+ captured_kwargs: dict[str, Any] = {}
class DummyLLMClient:
- async def __aenter__(self) -> "DummyLLMClient":
+ async def __aenter__(self) -> DummyLLMClient:
return self
async def __aexit__(self, exc_type, exc, tb) -> None:
@@ -731,7 +760,7 @@ async def run_eval(self, **kwargs):
captured_kwargs.update(kwargs)
return object()
- async def _verify_llm_client(*args, **kwargs) -> Optional[str]:
+ async def _verify_llm_client(*args, **kwargs) -> str | None:
return None
monkeypatch.setattr(
diff --git a/tests/test_rollout_api_key.py b/tests/test_rollout_api_key.py
index e14ae793..67408d94 100644
--- a/tests/test_rollout_api_key.py
+++ b/tests/test_rollout_api_key.py
@@ -16,12 +16,10 @@
from __future__ import annotations
-import pytest
-
from osmosis_ai.rollout.server.api_key import (
+ API_KEY_PREFIX,
generate_api_key,
validate_api_key,
- API_KEY_PREFIX,
)
@@ -42,7 +40,9 @@ def test_generate_api_key_is_url_safe(self) -> None:
"""Verify generated key contains only URL-safe characters."""
key = generate_api_key()
# URL-safe base64 uses only alphanumeric, '-', and '_'
- allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_")
+ allowed = set(
+ "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_"
+ )
for char in key:
assert char in allowed, f"Unexpected character: {char}"
diff --git a/tests/test_rollout_base.py b/tests/test_rollout_base.py
index 91521941..52e24f02 100644
--- a/tests/test_rollout_base.py
+++ b/tests/test_rollout_base.py
@@ -16,21 +16,19 @@
from __future__ import annotations
-from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock
import pytest
from osmosis_ai.rollout import (
+ OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
RolloutMetrics,
RolloutRequest,
RolloutResult,
- OpenAIFunctionToolSchema,
)
-
# =============================================================================
# RolloutResult Tests
# =============================================================================
@@ -343,7 +341,7 @@ async def test_agent_loop_run_can_return_complete_result(
class SimpleLoop(RolloutAgentLoop):
name = "simple"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -375,7 +373,7 @@ async def test_agent_loop_run_can_return_error_result(
class FailingLoop(RolloutAgentLoop):
name = "failing"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
diff --git a/tests/test_rollout_client.py b/tests/test_rollout_client.py
index 82f3ab86..fc4d76f3 100644
--- a/tests/test_rollout_client.py
+++ b/tests/test_rollout_client.py
@@ -17,7 +17,6 @@
from __future__ import annotations
import json
-from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@@ -33,7 +32,6 @@
)
from osmosis_ai.rollout.client import CompletionsResult
-
# =============================================================================
# OsmosisLLMClient Initialization Tests
# =============================================================================
@@ -765,7 +763,9 @@ async def test_client_async_context_manager() -> None:
"""Verify async context manager properly closes client."""
mock_http_client = AsyncMock()
# Patch httpx.AsyncClient to avoid creating a real SSL context in test env.
- with patch("osmosis_ai.rollout.client.httpx.AsyncClient", return_value=mock_http_client):
+ with patch(
+ "osmosis_ai.rollout.client.httpx.AsyncClient", return_value=mock_http_client
+ ):
async with OsmosisLLMClient(
server_url="http://localhost:8080",
rollout_id="test-123",
diff --git a/tests/test_rollout_registry.py b/tests/test_rollout_registry.py
index 3808d80f..3e9ca252 100644
--- a/tests/test_rollout_registry.py
+++ b/tests/test_rollout_registry.py
@@ -16,21 +16,19 @@
from __future__ import annotations
-from typing import List
-
import pytest
from osmosis_ai.rollout import (
AgentLoopNotFoundError,
+ OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
RolloutRequest,
RolloutResult,
- OpenAIFunctionToolSchema,
)
from osmosis_ai.rollout.registry import (
- AgentLoopRegistry,
_REGISTRY,
+ AgentLoopRegistry,
get_agent_loop,
list_agent_loops,
register_agent_loop,
@@ -43,7 +41,7 @@ class TestAgentLoop(RolloutAgentLoop):
name = "test_agent"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -55,7 +53,7 @@ class AnotherAgentLoop(RolloutAgentLoop):
name = "another_agent"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
diff --git a/tests/test_rollout_schemas.py b/tests/test_rollout_schemas.py
index 18d38009..fff9a68a 100644
--- a/tests/test_rollout_schemas.py
+++ b/tests/test_rollout_schemas.py
@@ -20,23 +20,22 @@
from pydantic import ValidationError
from osmosis_ai.rollout import (
+ DEFAULT_MAX_METADATA_SIZE_BYTES,
CompletionsRequest,
CompletionUsage,
- DEFAULT_MAX_METADATA_SIZE_BYTES,
InitResponse,
OpenAIFunctionParametersSchema,
OpenAIFunctionPropertySchema,
OpenAIFunctionSchema,
+ OpenAIFunctionToolSchema,
RolloutMetrics,
RolloutRequest,
RolloutResponse,
RolloutStatus,
- OpenAIFunctionToolSchema,
get_max_metadata_size_bytes,
set_max_metadata_size_bytes,
)
-
# =============================================================================
# RolloutRequest Tests
# =============================================================================
diff --git a/tests/test_rollout_server.py b/tests/test_rollout_server.py
index dec11b74..7974cc0a 100644
--- a/tests/test_rollout_server.py
+++ b/tests/test_rollout_server.py
@@ -17,17 +17,17 @@
from __future__ import annotations
import asyncio
-from typing import Any, Dict, List
+from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from osmosis_ai.rollout import (
+ OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
RolloutRequest,
RolloutResult,
- OpenAIFunctionToolSchema,
create_app,
)
from osmosis_ai.rollout.server import AppState
@@ -47,10 +47,10 @@ class SimpleAgentLoop(RolloutAgentLoop):
name = "simple_agent"
- def __init__(self, tools: List[OpenAIFunctionToolSchema] | None = None):
+ def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
self._tools = tools or []
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return self._tools
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -62,7 +62,7 @@ class FailingAgentLoop(RolloutAgentLoop):
name = "failing_agent"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -325,7 +325,9 @@ def test_init_endpoint_returns_202() -> None:
@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_returns_tools(sample_tool_schema: OpenAIFunctionToolSchema) -> None:
+def test_init_endpoint_returns_tools(
+ sample_tool_schema: OpenAIFunctionToolSchema,
+) -> None:
"""Verify /v1/rollout/init returns tools from agent_loop.get_tools()."""
agent_loop = SimpleAgentLoop(tools=[sample_tool_schema])
app = create_app(agent_loop)
@@ -476,7 +478,7 @@ class CountingToolsAgentLoop(RolloutAgentLoop):
def __init__(self) -> None:
self.calls = 0
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
self.calls += 1
return [
OpenAIFunctionToolSchema(
@@ -723,8 +725,8 @@ def test_create_app_raises_without_fastapi() -> None:
def test_rollout_context_debug_disabled_by_default() -> None:
"""Verify debug logging is disabled by default."""
- from osmosis_ai.rollout.core.base import RolloutContext
from osmosis_ai.rollout import RolloutRequest
+ from osmosis_ai.rollout.core.base import RolloutContext
request = RolloutRequest(
rollout_id="test-123",
@@ -746,8 +748,8 @@ def test_rollout_context_debug_disabled_by_default() -> None:
def test_rollout_context_debug_enabled_when_dir_set() -> None:
"""Verify debug logging is enabled when _debug_dir is set."""
- from osmosis_ai.rollout.core.base import RolloutContext
from osmosis_ai.rollout import RolloutRequest
+ from osmosis_ai.rollout.core.base import RolloutContext
request = RolloutRequest(
rollout_id="test-123",
@@ -770,8 +772,8 @@ def test_rollout_context_debug_enabled_when_dir_set() -> None:
def test_rollout_context_log_event_noop_when_disabled() -> None:
"""Verify log_event is a no-op when debug logging is disabled."""
- from osmosis_ai.rollout.core.base import RolloutContext
from osmosis_ai.rollout import RolloutRequest
+ from osmosis_ai.rollout.core.base import RolloutContext
request = RolloutRequest(
rollout_id="test-123",
@@ -794,8 +796,9 @@ def test_rollout_context_log_event_noop_when_disabled() -> None:
def test_rollout_context_log_event_writes_to_file(tmp_path: Any) -> None:
"""Verify log_event writes events to JSONL file."""
import json
- from osmosis_ai.rollout.core.base import RolloutContext
+
from osmosis_ai.rollout import RolloutRequest
+ from osmosis_ai.rollout.core.base import RolloutContext
debug_dir = str(tmp_path / "debug_logs")
@@ -821,11 +824,12 @@ def test_rollout_context_log_event_writes_to_file(tmp_path: Any) -> None:
# Verify file was created
import os
+
debug_file = os.path.join(debug_dir, "test-rollout-456.jsonl")
assert os.path.exists(debug_file)
# Verify contents
- with open(debug_file, "r", encoding="utf-8") as f:
+ with open(debug_file, encoding="utf-8") as f:
lines = f.readlines()
assert len(lines) == 3
@@ -848,8 +852,8 @@ def test_rollout_context_log_event_writes_to_file(tmp_path: Any) -> None:
def test_rollout_context_get_debug_file_path() -> None:
"""Verify _get_debug_file_path returns correct path."""
- from osmosis_ai.rollout.core.base import RolloutContext
from osmosis_ai.rollout import RolloutRequest
+ from osmosis_ai.rollout.core.base import RolloutContext
request = RolloutRequest(
rollout_id="my-rollout-id",
@@ -875,7 +879,9 @@ def test_rollout_context_get_debug_file_path() -> None:
llm=mock_llm,
_debug_dir="/var/log/rollouts",
)
- assert ctx_with_debug._get_debug_file_path() == "/var/log/rollouts/my-rollout-id.jsonl"
+ assert (
+ ctx_with_debug._get_debug_file_path() == "/var/log/rollouts/my-rollout-id.jsonl"
+ )
@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
@@ -894,6 +900,7 @@ async def test_background_rollout_with_debug_logging(tmp_path: Any) -> None:
"""Verify debug logging works in background rollout execution."""
import json
import os
+
from osmosis_ai.rollout import RolloutMetrics
debug_dir = str(tmp_path / "debug_logs")
@@ -903,7 +910,7 @@ class DebugLoggingAgentLoop(RolloutAgentLoop):
name = "debug_logging_agent"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -944,7 +951,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
assert os.path.exists(debug_file)
# Verify contents
- with open(debug_file, "r", encoding="utf-8") as f:
+ with open(debug_file, encoding="utf-8") as f:
lines = f.readlines()
assert len(lines) == 2
diff --git a/tests/test_rollout_testing.py b/tests/test_rollout_testing.py
index 1c81c7b2..daff4d58 100644
--- a/tests/test_rollout_testing.py
+++ b/tests/test_rollout_testing.py
@@ -28,14 +28,12 @@
# Check if FastAPI is available for testing
try:
from fastapi.testclient import TestClient
+
HAS_FASTAPI = True
except ImportError:
HAS_FASTAPI = False
-requires_fastapi = pytest.mark.skipif(
- not HAS_FASTAPI,
- reason="FastAPI not installed"
-)
+requires_fastapi = pytest.mark.skipif(not HAS_FASTAPI, reason="FastAPI not installed")
# =============================================================================
@@ -323,6 +321,7 @@ def test_mock_trainer_with_tracker() -> None:
@requires_fastapi
def test_mock_trainer_custom_tool_generator() -> None:
"""Verify mock trainer uses custom tool call generator."""
+
def custom_generator(message):
if "weather" in message.get("content", "").lower():
return [
diff --git a/tests/test_rollout_tools.py b/tests/test_rollout_tools.py
index fee6ba61..701efe9c 100644
--- a/tests/test_rollout_tools.py
+++ b/tests/test_rollout_tools.py
@@ -29,7 +29,6 @@
serialize_tool_result,
)
-
# =============================================================================
# create_tool_result Tests
# =============================================================================
@@ -203,7 +202,7 @@ def test_get_tool_call_info_missing_id() -> None:
"arguments": "{}",
},
}
- call_id, name, args = get_tool_call_info(tool_call)
+ call_id, _name, _args = get_tool_call_info(tool_call)
assert call_id == "unknown"
diff --git a/tests/test_rollout_utils.py b/tests/test_rollout_utils.py
index b08b51e6..da4d6f51 100644
--- a/tests/test_rollout_utils.py
+++ b/tests/test_rollout_utils.py
@@ -16,9 +16,7 @@
from __future__ import annotations
-from typing import Any, Dict
-
-import pytest
+from typing import Any
from osmosis_ai.rollout import (
count_messages_by_role,
@@ -31,14 +29,13 @@
parse_tool_calls,
)
-
# =============================================================================
# parse_tool_calls Tests
# =============================================================================
def test_parse_tool_calls_with_calls(
- sample_assistant_message_with_tool_calls: Dict[str, Any]
+ sample_assistant_message_with_tool_calls: dict[str, Any],
) -> None:
"""Verify parse_tool_calls extracts tool_calls."""
tool_calls = parse_tool_calls(sample_assistant_message_with_tool_calls)
@@ -46,7 +43,7 @@ def test_parse_tool_calls_with_calls(
assert tool_calls[0]["id"] == "call_123"
-def test_parse_tool_calls_no_calls(sample_assistant_message: Dict[str, Any]) -> None:
+def test_parse_tool_calls_no_calls(sample_assistant_message: dict[str, Any]) -> None:
"""Verify parse_tool_calls returns empty list when no tool_calls."""
tool_calls = parse_tool_calls(sample_assistant_message)
assert tool_calls == []
diff --git a/tests/test_rollout_validator.py b/tests/test_rollout_validator.py
index 8cf054a2..c441c393 100644
--- a/tests/test_rollout_validator.py
+++ b/tests/test_rollout_validator.py
@@ -16,7 +16,7 @@
from __future__ import annotations
-from typing import Any, Dict, List
+from typing import Any
import pytest
@@ -37,7 +37,6 @@
validate_agent_loop,
)
-
# =============================================================================
# Test Agent Loop Implementations
# =============================================================================
@@ -48,7 +47,7 @@ class ValidAgentLoop(RolloutAgentLoop):
name = "valid_agent"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return [
OpenAIFunctionToolSchema(
type="function",
@@ -76,7 +75,7 @@ class NoToolsAgentLoop(RolloutAgentLoop):
name = "no_tools"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
return []
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -88,7 +87,7 @@ class DictToolsAgentLoop(RolloutAgentLoop):
name = "dict_tools"
- def get_tools(self, request: RolloutRequest) -> List[Any]:
+ def get_tools(self, request: RolloutRequest) -> list[Any]:
return [
{
"type": "function",
@@ -116,7 +115,7 @@ class InvalidToolsAgentLoop(RolloutAgentLoop):
name = "invalid_tools"
- def get_tools(self, request: RolloutRequest) -> List[Any]:
+ def get_tools(self, request: RolloutRequest) -> list[Any]:
return [
{"invalid": "tool"}, # Missing type and function
]
@@ -130,7 +129,7 @@ class GetToolsRaisesAgentLoop(RolloutAgentLoop):
name = "get_tools_raises"
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
raise ValueError("Test exception in get_tools")
async def run(self, ctx: RolloutContext) -> RolloutResult:
@@ -154,7 +153,7 @@ class MissingFunctionNameAgentLoop(RolloutAgentLoop):
name = "missing_function_name"
- def get_tools(self, request: RolloutRequest) -> List[Any]:
+ def get_tools(self, request: RolloutRequest) -> list[Any]:
return [
{
"type": "function",
@@ -174,7 +173,7 @@ class MissingDescriptionAgentLoop(RolloutAgentLoop):
name = "missing_description"
- def get_tools(self, request: RolloutRequest) -> List[Any]:
+ def get_tools(self, request: RolloutRequest) -> list[Any]:
return [
{
"type": "function",
diff --git a/tests/test_test_mode_dataset.py b/tests/test_test_mode_dataset.py
index 85fa7462..0d2f7268 100644
--- a/tests/test_test_mode_dataset.py
+++ b/tests/test_test_mode_dataset.py
@@ -3,19 +3,19 @@
from __future__ import annotations
import json
-import tempfile
from pathlib import Path
-from typing import Any, Dict
+from typing import Any
import pytest
from osmosis_ai.rollout.eval.common.dataset import (
- REQUIRED_COLUMNS,
DatasetReader,
- DatasetRow,
dataset_row_to_request,
)
-from osmosis_ai.rollout.eval.common.errors import DatasetParseError, DatasetValidationError
+from osmosis_ai.rollout.eval.common.errors import (
+ DatasetParseError,
+ DatasetValidationError,
+)
# Check if pyarrow is available for Parquet tests
try:
@@ -355,7 +355,7 @@ class TestDatasetRowToRequest:
def test_basic_conversion(self) -> None:
"""Test basic conversion from DatasetRow to RolloutRequest."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "What is 2+2?",
"system_prompt": "You are a calculator.",
"ground_truth": "4",
@@ -373,7 +373,7 @@ def test_basic_conversion(self) -> None:
def test_ground_truth_in_metadata(self) -> None:
"""Test that ground_truth is stored in metadata."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Question",
"system_prompt": "System",
"ground_truth": "Expected Answer",
@@ -386,7 +386,7 @@ def test_ground_truth_in_metadata(self) -> None:
def test_extra_columns_in_metadata(self) -> None:
"""Test that extra columns are preserved in metadata."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Question",
"system_prompt": "System",
"ground_truth": "Answer",
@@ -401,7 +401,7 @@ def test_extra_columns_in_metadata(self) -> None:
def test_max_turns_parameter(self) -> None:
"""Test that max_turns is passed correctly."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Question",
"system_prompt": "System",
"ground_truth": "Answer",
@@ -413,7 +413,7 @@ def test_max_turns_parameter(self) -> None:
def test_completion_params(self) -> None:
"""Test that completion_params are passed correctly."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Question",
"system_prompt": "System",
"ground_truth": "Answer",
@@ -430,7 +430,7 @@ def test_completion_params(self) -> None:
def test_rollout_id_prefix_and_metadata_overrides(self) -> None:
"""Test custom rollout_id_prefix and metadata overrides."""
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Question",
"system_prompt": "System",
"ground_truth": "Answer",
diff --git a/tests/test_test_mode_providers.py b/tests/test_test_mode_providers.py
index b81f6e0c..3f3d4ad0 100644
--- a/tests/test_test_mode_providers.py
+++ b/tests/test_test_mode_providers.py
@@ -3,7 +3,7 @@
from __future__ import annotations
import types
-from typing import Any, Dict, List
+from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -17,9 +17,7 @@ class TestExternalLLMClient:
def test_model_auto_prefix(self) -> None:
"""Test that simple model names get auto-prefixed with openai/."""
- with patch(
- "osmosis_ai.rollout.eval.common.llm_client._get_litellm"
- ) as mock:
+ with patch("osmosis_ai.rollout.eval.common.llm_client._get_litellm") as mock:
mock.return_value = MagicMock()
from osmosis_ai.rollout.eval.common.llm_client import (
ExternalLLMClient,
@@ -35,9 +33,7 @@ def test_model_auto_prefix(self) -> None:
def test_set_and_clear_tools(self) -> None:
"""Test setting and clearing tools."""
- with patch(
- "osmosis_ai.rollout.eval.common.llm_client._get_litellm"
- ) as mock:
+ with patch("osmosis_ai.rollout.eval.common.llm_client._get_litellm") as mock:
mock.return_value = MagicMock()
from osmosis_ai.rollout.eval.common.llm_client import (
ExternalLLMClient,
@@ -76,9 +72,7 @@ def test_set_and_clear_tools(self) -> None:
def test_metrics_tracking(self) -> None:
"""Test metrics tracking methods."""
- with patch(
- "osmosis_ai.rollout.eval.common.llm_client._get_litellm"
- ) as mock:
+ with patch("osmosis_ai.rollout.eval.common.llm_client._get_litellm") as mock:
mock.return_value = MagicMock()
from osmosis_ai.rollout.eval.common.llm_client import (
ExternalLLMClient,
@@ -121,9 +115,7 @@ def test_metrics_tracking(self) -> None:
@pytest.mark.asyncio
async def test_context_manager(self) -> None:
"""Test async context manager protocol."""
- with patch(
- "osmosis_ai.rollout.eval.common.llm_client._get_litellm"
- ) as mock:
+ with patch("osmosis_ai.rollout.eval.common.llm_client._get_litellm") as mock:
mock.return_value = MagicMock()
from osmosis_ai.rollout.eval.common.llm_client import (
ExternalLLMClient,
@@ -193,9 +185,7 @@ async def test_chat_completions_calls_litellm(self) -> None:
)
client = ExternalLLMClient(model="gpt-4o")
- result = await client.chat_completions(
- [{"role": "user", "content": "Hi"}]
- )
+ result = await client.chat_completions([{"role": "user", "content": "Hi"}])
assert isinstance(result, CompletionsResult)
assert result.message["role"] == "assistant"
@@ -229,7 +219,7 @@ async def test_tools_auto_injection(self) -> None:
prompt_tokens=10, completion_tokens=5, total_tokens=15
)
- received_kwargs: Dict[str, Any] = {}
+ received_kwargs: dict[str, Any] = {}
async def capture_kwargs(**kwargs):
nonlocal received_kwargs
@@ -310,9 +300,7 @@ async def test_wraps_missing_provider_error_with_compact_hint(self) -> None:
)
with pytest.raises(SystemicProviderError) as exc_info:
- await client.chat_completions(
- [{"role": "user", "content": "hello"}]
- )
+ await client.chat_completions([{"role": "user", "content": "hello"}])
message = str(exc_info.value)
assert "Cannot connect to custom endpoint" in message
@@ -339,11 +327,17 @@ async def test_preflight_success(self) -> None:
mock_response = MagicMock()
mock_response.choices = [
MagicMock(
- message=MagicMock(model_dump=MagicMock(return_value={"role": "assistant", "content": "h"})),
+ message=MagicMock(
+ model_dump=MagicMock(
+ return_value={"role": "assistant", "content": "h"}
+ )
+ ),
finish_reason="stop",
)
]
- mock_response.usage = MagicMock(prompt_tokens=1, completion_tokens=1, total_tokens=2)
+ mock_response.usage = MagicMock(
+ prompt_tokens=1, completion_tokens=1, total_tokens=2
+ )
mock_litellm.acompletion = AsyncMock(return_value=mock_response)
with patch(
@@ -472,9 +466,10 @@ async def test_preflight_404_api_error(self) -> None:
with pytest.raises(SystemicProviderError):
await client.preflight_check()
-
@pytest.mark.asyncio
- async def test_preflight_connection_error_wrapped_as_internal_server_error(self) -> None:
+ async def test_preflight_connection_error_wrapped_as_internal_server_error(
+ self,
+ ) -> None:
"""Preflight catches connection errors wrapped as InternalServerError (APIError subclass)."""
mock_litellm = MagicMock()
for error_name in (
@@ -503,7 +498,9 @@ async def test_preflight_connection_error_wrapped_as_internal_server_error(self)
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
- client = ExternalLLMClient(model="gpt-4o", api_base="http://localhost:9999/v1")
+ client = ExternalLLMClient(
+ model="gpt-4o", api_base="http://localhost:9999/v1"
+ )
with pytest.raises(SystemicProviderError) as exc_info:
await client.preflight_check()
assert "Cannot connect" in str(exc_info.value)
@@ -595,7 +592,10 @@ async def test_rate_limit_is_not_systemic(self) -> None:
"osmosis_ai.rollout.eval.common.llm_client._get_litellm",
return_value=mock_litellm,
):
- from osmosis_ai.rollout.eval.common.errors import ProviderError, SystemicProviderError
+ from osmosis_ai.rollout.eval.common.errors import (
+ ProviderError,
+ SystemicProviderError,
+ )
from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
client = ExternalLLMClient(model="gpt-4o")
@@ -665,7 +665,9 @@ async def test_connection_error_wrapped_as_api_error_is_systemic(self) -> None:
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
- client = ExternalLLMClient(model="gpt-4o", api_base="http://localhost:9999/v1")
+ client = ExternalLLMClient(
+ model="gpt-4o", api_base="http://localhost:9999/v1"
+ )
with pytest.raises(SystemicProviderError) as exc_info:
await client.chat_completions([{"role": "user", "content": "hi"}])
assert "Cannot connect" in str(exc_info.value)
@@ -693,7 +695,10 @@ async def test_api_error_500_is_not_systemic(self) -> None:
"osmosis_ai.rollout.eval.common.llm_client._get_litellm",
return_value=mock_litellm,
):
- from osmosis_ai.rollout.eval.common.errors import ProviderError, SystemicProviderError
+ from osmosis_ai.rollout.eval.common.errors import (
+ ProviderError,
+ SystemicProviderError,
+ )
from osmosis_ai.rollout.eval.common.llm_client import ExternalLLMClient
client = ExternalLLMClient(model="gpt-4o")
@@ -710,17 +715,17 @@ def test_missing_litellm_raises_provider_error(self) -> None:
"""Test that missing LiteLLM raises ProviderError."""
from osmosis_ai.rollout.eval.common.errors import ProviderError
- with patch.dict("sys.modules", {"litellm": None}):
- with patch(
- "osmosis_ai.rollout.eval.common.llm_client._get_litellm"
- ) as mock:
- mock.side_effect = ProviderError("LiteLLM is required")
+ with (
+ patch.dict("sys.modules", {"litellm": None}),
+ patch("osmosis_ai.rollout.eval.common.llm_client._get_litellm") as mock,
+ ):
+ mock.side_effect = ProviderError("LiteLLM is required")
- with pytest.raises(ProviderError) as exc_info:
- from osmosis_ai.rollout.eval.common.llm_client import (
- ExternalLLMClient,
- )
+ with pytest.raises(ProviderError) as exc_info:
+ from osmosis_ai.rollout.eval.common.llm_client import (
+ ExternalLLMClient,
+ )
- ExternalLLMClient()
+ ExternalLLMClient()
- assert "LiteLLM" in str(exc_info.value)
+ assert "LiteLLM" in str(exc_info.value)
diff --git a/tests/test_test_mode_runner.py b/tests/test_test_mode_runner.py
index 70bd07b1..32909d3c 100644
--- a/tests/test_test_mode_runner.py
+++ b/tests/test_test_mode_runner.py
@@ -2,8 +2,8 @@
from __future__ import annotations
-from typing import Any, Dict, List
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing import Any
+from unittest.mock import MagicMock, patch
import pytest
@@ -21,8 +21,8 @@
from osmosis_ai.rollout.eval.common.dataset import DatasetRow
from osmosis_ai.rollout.eval.test_mode.runner import (
LocalTestBatchResult,
- LocalTestRunResult,
LocalTestRunner,
+ LocalTestRunResult,
)
@@ -33,12 +33,12 @@ class MockTestLLMClient:
"""
def __init__(self) -> None:
- self._tools: List[Dict[str, Any]] | None = None
+ self._tools: list[dict[str, Any]] | None = None
self._llm_latency_ms: float = 0.0
self._num_llm_calls: int = 0
self._prompt_tokens: int = 0
self._response_tokens: int = 0
- self.completions_calls: List[Dict[str, Any]] = []
+ self.completions_calls: list[dict[str, Any]] = []
self.mock_response = CompletionsResult(
message={"role": "assistant", "content": "Test response"},
token_ids=[],
@@ -47,7 +47,7 @@ def __init__(self) -> None:
finish_reason="stop",
)
- def set_tools(self, tools: List[Any]) -> None:
+ def set_tools(self, tools: list[Any]) -> None:
if tools:
self._tools = [
t.model_dump(exclude_none=True) if hasattr(t, "model_dump") else t
@@ -82,7 +82,7 @@ def _record_usage(
self._response_tokens += completion_tokens
async def chat_completions(
- self, messages: List[Dict[str, Any]], **kwargs: Any
+ self, messages: list[dict[str, Any]], **kwargs: Any
) -> CompletionsResult:
# Auto-inject tools if set
if self._tools is not None and "tools" not in kwargs:
@@ -99,7 +99,7 @@ class MockAgentLoop(RolloutAgentLoop):
def __init__(
self,
- tools: List[OpenAIFunctionToolSchema] | None = None,
+ tools: list[OpenAIFunctionToolSchema] | None = None,
run_result: RolloutResult | None = None,
run_error: Exception | None = None,
call_llm: bool = False,
@@ -108,10 +108,10 @@ def __init__(
self._run_result = run_result
self._run_error = run_error
self._call_llm = call_llm
- self.get_tools_calls: List[RolloutRequest] = []
- self.run_calls: List[RolloutContext] = []
+ self.get_tools_calls: list[RolloutRequest] = []
+ self.run_calls: list[RolloutContext] = []
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
self.get_tools_calls.append(request)
return self._tools
@@ -334,7 +334,7 @@ async def test_run_batch_progress_callback(self) -> None:
agent = MockAgentLoop(tools=[create_sample_tool()])
runner = LocalTestRunner(agent_loop=agent, llm_client=client)
- progress_calls: List[tuple] = []
+ progress_calls: list[tuple] = []
def on_progress(current: int, total: int, result: LocalTestRunResult) -> None:
progress_calls.append((current, total, result))
@@ -380,7 +380,7 @@ async def test_run_batch_with_start_index_and_progress(self) -> None:
agent = MockAgentLoop(tools=[create_sample_tool()])
runner = LocalTestRunner(agent_loop=agent, llm_client=client)
- progress_calls: List[tuple] = []
+ progress_calls: list[tuple] = []
def on_progress(current: int, total: int, result: LocalTestRunResult) -> None:
progress_calls.append((current, total, result.row_index))
@@ -456,7 +456,7 @@ async def run_single(
row: DatasetRow,
row_index: int,
max_turns: int = 10,
- completion_params: Dict[str, Any] | None = None,
+ completion_params: dict[str, Any] | None = None,
) -> LocalTestRunResult:
self._calls += 1
if self._calls == 1:
@@ -486,6 +486,7 @@ async def run_single(
@pytest.mark.asyncio
async def test_run_batch_normal_errors_dont_stop_early(self) -> None:
"""Regular exceptions should not stop the batch early."""
+
class FailingAgent(MockAgentLoop):
async def run(self, ctx: RolloutContext) -> RolloutResult:
raise RuntimeError("Transient error")
@@ -634,7 +635,7 @@ async def test_context_has_correct_request(self) -> None:
agent = MockAgentLoop(tools=[])
runner = LocalTestRunner(agent_loop=agent, llm_client=client)
- row: Dict[str, Any] = {
+ row: dict[str, Any] = {
"user_prompt": "Test question",
"system_prompt": "Test system",
"ground_truth": "Test answer",
diff --git a/uv.lock b/uv.lock
index 9d8085aa..67e1f85d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -198,6 +198,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" },
]
+[[package]]
+name = "authlib"
+version = "1.6.8"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cryptography" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6b/6c/c88eac87468c607f88bc24df1f3b31445ee6fc9ba123b09e666adf687cd9/authlib-1.6.8.tar.gz", hash = "sha256:41ae180a17cf672bc784e4a518e5c82687f1fe1e98b0cafaeda80c8e4ab2d1cb", size = 165074, upload-time = "2026-02-14T04:02:17.941Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/9b/73/f7084bf12755113cd535ae586782ff3a6e710bfbe6a0d13d1c2f81ffbbfa/authlib-1.6.8-py2.py3-none-any.whl", hash = "sha256:97286fd7a15e6cfefc32771c8ef9c54f0ed58028f1322de6a2a7c969c3817888", size = 244116, upload-time = "2026-02-14T04:02:15.579Z" },
+]
+
[[package]]
name = "backports-asyncio-runner"
version = "1.2.0"
@@ -208,47 +220,30 @@ wheels = [
]
[[package]]
-name = "black"
-version = "25.12.0"
+name = "backports-tarfile"
+version = "1.2.0"
source = { registry = "https://pypi.org/simple" }
-dependencies = [
- { name = "click" },
- { name = "mypy-extensions" },
- { name = "packaging" },
- { name = "pathspec" },
- { name = "platformdirs" },
- { name = "pytokens" },
- { name = "tomli", marker = "python_full_version < '3.11'" },
- { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+sdist = { url = "https://files.pythonhosted.org/packages/86/72/cd9b395f25e290e633655a100af28cb253e4393396264a98bd5f5951d50f/backports_tarfile-1.2.0.tar.gz", hash = "sha256:d75e02c268746e1b8144c278978b6e98e85de6ad16f8e4b0844a154557eca991", size = 86406, upload-time = "2024-05-28T17:01:54.731Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b9/fa/123043af240e49752f1c4bd24da5053b6bd00cad78c2be53c0d1e8b975bc/backports.tarfile-1.2.0-py3-none-any.whl", hash = "sha256:77e284d754527b01fb1e6fa8a1afe577858ebe4e9dad8919e34c862cb399bc34", size = 30181, upload-time = "2024-05-28T17:01:53.112Z" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/c4/d9/07b458a3f1c525ac392b5edc6b191ff140b596f9d77092429417a54e249d/black-25.12.0.tar.gz", hash = "sha256:8d3dd9cea14bff7ddc0eb243c811cdb1a011ebb4800a5f0335a01a68654796a7", size = 659264, upload-time = "2025-12-08T01:40:52.501Z" }
-wheels = [
- { url = "https://files.pythonhosted.org/packages/37/d5/8d3145999d380e5d09bb00b0f7024bf0a8ccb5c07b5648e9295f02ec1d98/black-25.12.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f85ba1ad15d446756b4ab5f3044731bf68b777f8f9ac9cdabd2425b97cd9c4e8", size = 1895720, upload-time = "2025-12-08T01:46:58.197Z" },
- { url = "https://files.pythonhosted.org/packages/06/97/7acc85c4add41098f4f076b21e3e4e383ad6ed0a3da26b2c89627241fc11/black-25.12.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:546eecfe9a3a6b46f9d69d8a642585a6eaf348bcbbc4d87a19635570e02d9f4a", size = 1727193, upload-time = "2025-12-08T01:52:26.674Z" },
- { url = "https://files.pythonhosted.org/packages/24/f0/fdf0eb8ba907ddeb62255227d29d349e8256ef03558fbcadfbc26ecfe3b2/black-25.12.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:17dcc893da8d73d8f74a596f64b7c98ef5239c2cd2b053c0f25912c4494bf9ea", size = 1774506, upload-time = "2025-12-08T01:46:25.721Z" },
- { url = "https://files.pythonhosted.org/packages/e4/f5/9203a78efe00d13336786b133c6180a9303d46908a9aa72d1104ca214222/black-25.12.0-cp310-cp310-win_amd64.whl", hash = "sha256:09524b0e6af8ba7a3ffabdfc7a9922fb9adef60fed008c7cd2fc01f3048e6e6f", size = 1416085, upload-time = "2025-12-08T01:46:06.073Z" },
- { url = "https://files.pythonhosted.org/packages/ba/cc/7a6090e6b081c3316282c05c546e76affdce7bf7a3b7d2c3a2a69438bd01/black-25.12.0-cp310-cp310-win_arm64.whl", hash = "sha256:b162653ed89eb942758efeb29d5e333ca5bb90e5130216f8369857db5955a7da", size = 1226038, upload-time = "2025-12-08T01:45:29.388Z" },
- { url = "https://files.pythonhosted.org/packages/60/ad/7ac0d0e1e0612788dbc48e62aef8a8e8feffac7eb3d787db4e43b8462fa8/black-25.12.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d0cfa263e85caea2cff57d8f917f9f51adae8e20b610e2b23de35b5b11ce691a", size = 1877003, upload-time = "2025-12-08T01:43:29.967Z" },
- { url = "https://files.pythonhosted.org/packages/e8/dd/a237e9f565f3617a88b49284b59cbca2a4f56ebe68676c1aad0ce36a54a7/black-25.12.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1a2f578ae20c19c50a382286ba78bfbeafdf788579b053d8e4980afb079ab9be", size = 1712639, upload-time = "2025-12-08T01:52:46.756Z" },
- { url = "https://files.pythonhosted.org/packages/12/80/e187079df1ea4c12a0c63282ddd8b81d5107db6d642f7d7b75a6bcd6fc21/black-25.12.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3e1b65634b0e471d07ff86ec338819e2ef860689859ef4501ab7ac290431f9b", size = 1758143, upload-time = "2025-12-08T01:45:29.137Z" },
- { url = "https://files.pythonhosted.org/packages/93/b5/3096ccee4f29dc2c3aac57274326c4d2d929a77e629f695f544e159bfae4/black-25.12.0-cp311-cp311-win_amd64.whl", hash = "sha256:a3fa71e3b8dd9f7c6ac4d818345237dfb4175ed3bf37cd5a581dbc4c034f1ec5", size = 1420698, upload-time = "2025-12-08T01:45:53.379Z" },
- { url = "https://files.pythonhosted.org/packages/7e/39/f81c0ffbc25ffbe61c7d0385bf277e62ffc3e52f5ee668d7369d9854fadf/black-25.12.0-cp311-cp311-win_arm64.whl", hash = "sha256:51e267458f7e650afed8445dc7edb3187143003d52a1b710c7321aef22aa9655", size = 1229317, upload-time = "2025-12-08T01:46:35.606Z" },
- { url = "https://files.pythonhosted.org/packages/d1/bd/26083f805115db17fda9877b3c7321d08c647df39d0df4c4ca8f8450593e/black-25.12.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:31f96b7c98c1ddaeb07dc0f56c652e25bdedaac76d5b68a059d998b57c55594a", size = 1924178, upload-time = "2025-12-08T01:49:51.048Z" },
- { url = "https://files.pythonhosted.org/packages/89/6b/ea00d6651561e2bdd9231c4177f4f2ae19cc13a0b0574f47602a7519b6ca/black-25.12.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:05dd459a19e218078a1f98178c13f861fe6a9a5f88fc969ca4d9b49eb1809783", size = 1742643, upload-time = "2025-12-08T01:49:59.09Z" },
- { url = "https://files.pythonhosted.org/packages/6d/f3/360fa4182e36e9875fabcf3a9717db9d27a8d11870f21cff97725c54f35b/black-25.12.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1f68c5eff61f226934be6b5b80296cf6939e5d2f0c2f7d543ea08b204bfaf59", size = 1800158, upload-time = "2025-12-08T01:44:27.301Z" },
- { url = "https://files.pythonhosted.org/packages/f8/08/2c64830cb6616278067e040acca21d4f79727b23077633953081c9445d61/black-25.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:274f940c147ddab4442d316b27f9e332ca586d39c85ecf59ebdea82cc9ee8892", size = 1426197, upload-time = "2025-12-08T01:45:51.198Z" },
- { url = "https://files.pythonhosted.org/packages/d4/60/a93f55fd9b9816b7432cf6842f0e3000fdd5b7869492a04b9011a133ee37/black-25.12.0-cp312-cp312-win_arm64.whl", hash = "sha256:169506ba91ef21e2e0591563deda7f00030cb466e747c4b09cb0a9dae5db2f43", size = 1237266, upload-time = "2025-12-08T01:45:10.556Z" },
- { url = "https://files.pythonhosted.org/packages/c8/52/c551e36bc95495d2aa1a37d50566267aa47608c81a53f91daa809e03293f/black-25.12.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a05ddeb656534c3e27a05a29196c962877c83fa5503db89e68857d1161ad08a5", size = 1923809, upload-time = "2025-12-08T01:46:55.126Z" },
- { url = "https://files.pythonhosted.org/packages/a0/f7/aac9b014140ee56d247e707af8db0aae2e9efc28d4a8aba92d0abd7ae9d1/black-25.12.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9ec77439ef3e34896995503865a85732c94396edcc739f302c5673a2315e1e7f", size = 1742384, upload-time = "2025-12-08T01:49:37.022Z" },
- { url = "https://files.pythonhosted.org/packages/74/98/38aaa018b2ab06a863974c12b14a6266badc192b20603a81b738c47e902e/black-25.12.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e509c858adf63aa61d908061b52e580c40eae0dfa72415fa47ac01b12e29baf", size = 1798761, upload-time = "2025-12-08T01:46:05.386Z" },
- { url = "https://files.pythonhosted.org/packages/16/3a/a8ac542125f61574a3f015b521ca83b47321ed19bb63fe6d7560f348bfe1/black-25.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:252678f07f5bac4ff0d0e9b261fbb029fa530cfa206d0a636a34ab445ef8ca9d", size = 1429180, upload-time = "2025-12-08T01:45:34.903Z" },
- { url = "https://files.pythonhosted.org/packages/e6/2d/bdc466a3db9145e946762d52cd55b1385509d9f9004fec1c97bdc8debbfb/black-25.12.0-cp313-cp313-win_arm64.whl", hash = "sha256:bc5b1c09fe3c931ddd20ee548511c64ebf964ada7e6f0763d443947fd1c603ce", size = 1239350, upload-time = "2025-12-08T01:46:09.458Z" },
- { url = "https://files.pythonhosted.org/packages/35/46/1d8f2542210c502e2ae1060b2e09e47af6a5e5963cb78e22ec1a11170b28/black-25.12.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:0a0953b134f9335c2434864a643c842c44fba562155c738a2a37a4d61f00cad5", size = 1917015, upload-time = "2025-12-08T01:53:27.987Z" },
- { url = "https://files.pythonhosted.org/packages/41/37/68accadf977672beb8e2c64e080f568c74159c1aaa6414b4cd2aef2d7906/black-25.12.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2355bbb6c3b76062870942d8cc450d4f8ac71f9c93c40122762c8784df49543f", size = 1741830, upload-time = "2025-12-08T01:54:36.861Z" },
- { url = "https://files.pythonhosted.org/packages/ac/76/03608a9d8f0faad47a3af3a3c8c53af3367f6c0dd2d23a84710456c7ac56/black-25.12.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9678bd991cc793e81d19aeeae57966ee02909877cb65838ccffef24c3ebac08f", size = 1791450, upload-time = "2025-12-08T01:44:52.581Z" },
- { url = "https://files.pythonhosted.org/packages/06/99/b2a4bd7dfaea7964974f947e1c76d6886d65fe5d24f687df2d85406b2609/black-25.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:97596189949a8aad13ad12fcbb4ae89330039b96ad6742e6f6b45e75ad5cfd83", size = 1452042, upload-time = "2025-12-08T01:46:13.188Z" },
- { url = "https://files.pythonhosted.org/packages/b2/7c/d9825de75ae5dd7795d007681b752275ea85a1c5d83269b4b9c754c2aaab/black-25.12.0-cp314-cp314-win_arm64.whl", hash = "sha256:778285d9ea197f34704e3791ea9404cd6d07595745907dd2ce3da7a13627b29b", size = 1267446, upload-time = "2025-12-08T01:46:14.497Z" },
- { url = "https://files.pythonhosted.org/packages/68/11/21331aed19145a952ad28fca2756a1433ee9308079bd03bd898e903a2e53/black-25.12.0-py3-none-any.whl", hash = "sha256:48ceb36c16dbc84062740049eef990bb2ce07598272e673c17d1a7720c71c828", size = 206191, upload-time = "2025-12-08T01:40:50.963Z" },
+
+[[package]]
+name = "beartype"
+version = "0.22.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c7/94/1009e248bbfbab11397abca7193bea6626806be9a327d399810d523a07cb/beartype-0.22.9.tar.gz", hash = "sha256:8f82b54aa723a2848a56008d18875f91c1db02c32ef6a62319a002e3e25a975f", size = 1608866, upload-time = "2025-12-13T06:50:30.72Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl", hash = "sha256:d16c9bbc61ea14637596c5f6fbff2ee99cbe3573e46a716401734ef50c3060c2", size = 1333658, upload-time = "2025-12-13T06:50:28.266Z" },
+]
+
+[[package]]
+name = "cachetools"
+version = "7.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d4/07/56595285564e90777d758ebd383d6b0b971b87729bbe2184a849932a3736/cachetools-7.0.1.tar.gz", hash = "sha256:e31e579d2c5b6e2944177a0397150d312888ddf4e16e12f1016068f0c03b8341", size = 36126, upload-time = "2026-02-10T22:24:05.03Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ed/9e/5faefbf9db1db466d633735faceda1f94aa99ce506ac450d232536266b32/cachetools-7.0.1-py3-none-any.whl", hash = "sha256:8f086515c254d5664ae2146d14fc7f65c9a4bce75152eb247e5a9c5e6d7b2ecf", size = 13484, upload-time = "2026-02-10T22:24:03.741Z" },
]
[[package]]
@@ -260,6 +255,88 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/70/7d/9bc192684cea499815ff478dfcdc13835ddf401365057044fb721ec6bddb/certifi-2025.11.12-py3-none-any.whl", hash = "sha256:97de8790030bbd5c2d96b7ec782fc2f7820ef8dba6db909ccf95449f2d062d4b", size = 159438, upload-time = "2025-11-12T02:54:49.735Z" },
]
+[[package]]
+name = "cffi"
+version = "2.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "pycparser", marker = "implementation_name != 'PyPy'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/93/d7/516d984057745a6cd96575eea814fe1edd6646ee6efd552fb7b0921dec83/cffi-2.0.0-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:0cf2d91ecc3fcc0625c2c530fe004f82c110405f101548512cce44322fa8ac44", size = 184283, upload-time = "2025-09-08T23:22:08.01Z" },
+ { url = "https://files.pythonhosted.org/packages/9e/84/ad6a0b408daa859246f57c03efd28e5dd1b33c21737c2db84cae8c237aa5/cffi-2.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f73b96c41e3b2adedc34a7356e64c8eb96e03a3782b535e043a986276ce12a49", size = 180504, upload-time = "2025-09-08T23:22:10.637Z" },
+ { url = "https://files.pythonhosted.org/packages/50/bd/b1a6362b80628111e6653c961f987faa55262b4002fcec42308cad1db680/cffi-2.0.0-cp310-cp310-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:53f77cbe57044e88bbd5ed26ac1d0514d2acf0591dd6bb02a3ae37f76811b80c", size = 208811, upload-time = "2025-09-08T23:22:12.267Z" },
+ { url = "https://files.pythonhosted.org/packages/4f/27/6933a8b2562d7bd1fb595074cf99cc81fc3789f6a6c05cdabb46284a3188/cffi-2.0.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3e837e369566884707ddaf85fc1744b47575005c0a229de3327f8f9a20f4efeb", size = 216402, upload-time = "2025-09-08T23:22:13.455Z" },
+ { url = "https://files.pythonhosted.org/packages/05/eb/b86f2a2645b62adcfff53b0dd97e8dfafb5c8aa864bd0d9a2c2049a0d551/cffi-2.0.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5eda85d6d1879e692d546a078b44251cdd08dd1cfb98dfb77b670c97cee49ea0", size = 203217, upload-time = "2025-09-08T23:22:14.596Z" },
+ { url = "https://files.pythonhosted.org/packages/9f/e0/6cbe77a53acf5acc7c08cc186c9928864bd7c005f9efd0d126884858a5fe/cffi-2.0.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9332088d75dc3241c702d852d4671613136d90fa6881da7d770a483fd05248b4", size = 203079, upload-time = "2025-09-08T23:22:15.769Z" },
+ { url = "https://files.pythonhosted.org/packages/98/29/9b366e70e243eb3d14a5cb488dfd3a0b6b2f1fb001a203f653b93ccfac88/cffi-2.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc7de24befaeae77ba923797c7c87834c73648a05a4bde34b3b7e5588973a453", size = 216475, upload-time = "2025-09-08T23:22:17.427Z" },
+ { url = "https://files.pythonhosted.org/packages/21/7a/13b24e70d2f90a322f2900c5d8e1f14fa7e2a6b3332b7309ba7b2ba51a5a/cffi-2.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cf364028c016c03078a23b503f02058f1814320a56ad535686f90565636a9495", size = 218829, upload-time = "2025-09-08T23:22:19.069Z" },
+ { url = "https://files.pythonhosted.org/packages/60/99/c9dc110974c59cc981b1f5b66e1d8af8af764e00f0293266824d9c4254bc/cffi-2.0.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e11e82b744887154b182fd3e7e8512418446501191994dbf9c9fc1f32cc8efd5", size = 211211, upload-time = "2025-09-08T23:22:20.588Z" },
+ { url = "https://files.pythonhosted.org/packages/49/72/ff2d12dbf21aca1b32a40ed792ee6b40f6dc3a9cf1644bd7ef6e95e0ac5e/cffi-2.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8ea985900c5c95ce9db1745f7933eeef5d314f0565b27625d9a10ec9881e1bfb", size = 218036, upload-time = "2025-09-08T23:22:22.143Z" },
+ { url = "https://files.pythonhosted.org/packages/e2/cc/027d7fb82e58c48ea717149b03bcadcbdc293553edb283af792bd4bcbb3f/cffi-2.0.0-cp310-cp310-win32.whl", hash = "sha256:1f72fb8906754ac8a2cc3f9f5aaa298070652a0ffae577e0ea9bd480dc3c931a", size = 172184, upload-time = "2025-09-08T23:22:23.328Z" },
+ { url = "https://files.pythonhosted.org/packages/33/fa/072dd15ae27fbb4e06b437eb6e944e75b068deb09e2a2826039e49ee2045/cffi-2.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:b18a3ed7d5b3bd8d9ef7a8cb226502c6bf8308df1525e1cc676c3680e7176739", size = 182790, upload-time = "2025-09-08T23:22:24.752Z" },
+ { url = "https://files.pythonhosted.org/packages/12/4a/3dfd5f7850cbf0d06dc84ba9aa00db766b52ca38d8b86e3a38314d52498c/cffi-2.0.0-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:b4c854ef3adc177950a8dfc81a86f5115d2abd545751a304c5bcf2c2c7283cfe", size = 184344, upload-time = "2025-09-08T23:22:26.456Z" },
+ { url = "https://files.pythonhosted.org/packages/4f/8b/f0e4c441227ba756aafbe78f117485b25bb26b1c059d01f137fa6d14896b/cffi-2.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2de9a304e27f7596cd03d16f1b7c72219bd944e99cc52b84d0145aefb07cbd3c", size = 180560, upload-time = "2025-09-08T23:22:28.197Z" },
+ { url = "https://files.pythonhosted.org/packages/b1/b7/1200d354378ef52ec227395d95c2576330fd22a869f7a70e88e1447eb234/cffi-2.0.0-cp311-cp311-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:baf5215e0ab74c16e2dd324e8ec067ef59e41125d3eade2b863d294fd5035c92", size = 209613, upload-time = "2025-09-08T23:22:29.475Z" },
+ { url = "https://files.pythonhosted.org/packages/b8/56/6033f5e86e8cc9bb629f0077ba71679508bdf54a9a5e112a3c0b91870332/cffi-2.0.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:730cacb21e1bdff3ce90babf007d0a0917cc3e6492f336c2f0134101e0944f93", size = 216476, upload-time = "2025-09-08T23:22:31.063Z" },
+ { url = "https://files.pythonhosted.org/packages/dc/7f/55fecd70f7ece178db2f26128ec41430d8720f2d12ca97bf8f0a628207d5/cffi-2.0.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:6824f87845e3396029f3820c206e459ccc91760e8fa24422f8b0c3d1731cbec5", size = 203374, upload-time = "2025-09-08T23:22:32.507Z" },
+ { url = "https://files.pythonhosted.org/packages/84/ef/a7b77c8bdc0f77adc3b46888f1ad54be8f3b7821697a7b89126e829e676a/cffi-2.0.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9de40a7b0323d889cf8d23d1ef214f565ab154443c42737dfe52ff82cf857664", size = 202597, upload-time = "2025-09-08T23:22:34.132Z" },
+ { url = "https://files.pythonhosted.org/packages/d7/91/500d892b2bf36529a75b77958edfcd5ad8e2ce4064ce2ecfeab2125d72d1/cffi-2.0.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8941aaadaf67246224cee8c3803777eed332a19d909b47e29c9842ef1e79ac26", size = 215574, upload-time = "2025-09-08T23:22:35.443Z" },
+ { url = "https://files.pythonhosted.org/packages/44/64/58f6255b62b101093d5df22dcb752596066c7e89dd725e0afaed242a61be/cffi-2.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:a05d0c237b3349096d3981b727493e22147f934b20f6f125a3eba8f994bec4a9", size = 218971, upload-time = "2025-09-08T23:22:36.805Z" },
+ { url = "https://files.pythonhosted.org/packages/ab/49/fa72cebe2fd8a55fbe14956f9970fe8eb1ac59e5df042f603ef7c8ba0adc/cffi-2.0.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:94698a9c5f91f9d138526b48fe26a199609544591f859c870d477351dc7b2414", size = 211972, upload-time = "2025-09-08T23:22:38.436Z" },
+ { url = "https://files.pythonhosted.org/packages/0b/28/dd0967a76aab36731b6ebfe64dec4e981aff7e0608f60c2d46b46982607d/cffi-2.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5fed36fccc0612a53f1d4d9a816b50a36702c28a2aa880cb8a122b3466638743", size = 217078, upload-time = "2025-09-08T23:22:39.776Z" },
+ { url = "https://files.pythonhosted.org/packages/2b/c0/015b25184413d7ab0a410775fdb4a50fca20f5589b5dab1dbbfa3baad8ce/cffi-2.0.0-cp311-cp311-win32.whl", hash = "sha256:c649e3a33450ec82378822b3dad03cc228b8f5963c0c12fc3b1e0ab940f768a5", size = 172076, upload-time = "2025-09-08T23:22:40.95Z" },
+ { url = "https://files.pythonhosted.org/packages/ae/8f/dc5531155e7070361eb1b7e4c1a9d896d0cb21c49f807a6c03fd63fc877e/cffi-2.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:66f011380d0e49ed280c789fbd08ff0d40968ee7b665575489afa95c98196ab5", size = 182820, upload-time = "2025-09-08T23:22:42.463Z" },
+ { url = "https://files.pythonhosted.org/packages/95/5c/1b493356429f9aecfd56bc171285a4c4ac8697f76e9bbbbb105e537853a1/cffi-2.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:c6638687455baf640e37344fe26d37c404db8b80d037c3d29f58fe8d1c3b194d", size = 177635, upload-time = "2025-09-08T23:22:43.623Z" },
+ { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" },
+ { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" },
+ { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" },
+ { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" },
+ { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" },
+ { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" },
+ { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" },
+ { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" },
+ { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" },
+ { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" },
+ { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" },
+ { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" },
+ { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" },
+ { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" },
+ { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" },
+ { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" },
+ { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" },
+ { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" },
+ { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" },
+ { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" },
+ { url = "https://files.pythonhosted.org/packages/92/c4/3ce07396253a83250ee98564f8d7e9789fab8e58858f35d07a9a2c78de9f/cffi-2.0.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fc33c5141b55ed366cfaad382df24fe7dcbc686de5be719b207bb248e3053dc5", size = 185320, upload-time = "2025-09-08T23:23:18.087Z" },
+ { url = "https://files.pythonhosted.org/packages/59/dd/27e9fa567a23931c838c6b02d0764611c62290062a6d4e8ff7863daf9730/cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c654de545946e0db659b3400168c9ad31b5d29593291482c43e3564effbcee13", size = 181487, upload-time = "2025-09-08T23:23:19.622Z" },
+ { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" },
+ { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" },
+ { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" },
+ { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" },
+ { url = "https://files.pythonhosted.org/packages/3e/aa/df335faa45b395396fcbc03de2dfcab242cd61a9900e914fe682a59170b1/cffi-2.0.0-cp314-cp314-win32.whl", hash = "sha256:087067fa8953339c723661eda6b54bc98c5625757ea62e95eb4898ad5e776e9f", size = 175328, upload-time = "2025-09-08T23:23:44.61Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/92/882c2d30831744296ce713f0feb4c1cd30f346ef747b530b5318715cc367/cffi-2.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:203a48d1fb583fc7d78a4c6655692963b860a417c0528492a6bc21f1aaefab25", size = 185650, upload-time = "2025-09-08T23:23:45.848Z" },
+ { url = "https://files.pythonhosted.org/packages/9f/2c/98ece204b9d35a7366b5b2c6539c350313ca13932143e79dc133ba757104/cffi-2.0.0-cp314-cp314-win_arm64.whl", hash = "sha256:dbd5c7a25a7cb98f5ca55d258b103a2054f859a46ae11aaf23134f9cc0d356ad", size = 180687, upload-time = "2025-09-08T23:23:47.105Z" },
+ { url = "https://files.pythonhosted.org/packages/3e/61/c768e4d548bfa607abcda77423448df8c471f25dbe64fb2ef6d555eae006/cffi-2.0.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:9a67fc9e8eb39039280526379fb3a70023d77caec1852002b4da7e8b270c4dd9", size = 188773, upload-time = "2025-09-08T23:23:29.347Z" },
+ { url = "https://files.pythonhosted.org/packages/2c/ea/5f76bce7cf6fcd0ab1a1058b5af899bfbef198bea4d5686da88471ea0336/cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7a66c7204d8869299919db4d5069a82f1561581af12b11b3c9f48c584eb8743d", size = 185013, upload-time = "2025-09-08T23:23:30.63Z" },
+ { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" },
+ { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" },
+ { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" },
+ { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" },
+ { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" },
+ { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" },
+ { url = "https://files.pythonhosted.org/packages/a0/1d/ec1a60bd1a10daa292d3cd6bb0b359a81607154fb8165f3ec95fe003b85c/cffi-2.0.0-cp314-cp314t-win32.whl", hash = "sha256:1fc9ea04857caf665289b7a75923f2c6ed559b8298a1b8c49e59f7dd95c8481e", size = 180487, upload-time = "2025-09-08T23:23:40.423Z" },
+ { url = "https://files.pythonhosted.org/packages/bf/41/4c1168c74fac325c0c8156f04b6749c8b6a8f405bbf91413ba088359f60d/cffi-2.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d68b6cef7827e8641e8ef16f4494edda8b36104d79773a334beaa1e3521430f6", size = 191726, upload-time = "2025-09-08T23:23:41.742Z" },
+ { url = "https://files.pythonhosted.org/packages/ae/3a/dbeec9d1ee0844c679f6bb5d6ad4e9f198b1224f4e7a32825f47f6192b0c/cffi-2.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0a1527a803f0a659de1af2e1fd700213caba79377e27e4693648c2923da066f9", size = 184195, upload-time = "2025-09-08T23:23:43.004Z" },
+]
+
[[package]]
name = "charset-normalizer"
version = "3.4.4"
@@ -361,6 +438,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" },
]
+[[package]]
+name = "cloudpickle"
+version = "3.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" },
+]
+
[[package]]
name = "colorama"
version = "0.4.6"
@@ -370,6 +456,105 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
]
+[[package]]
+name = "croniter"
+version = "6.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "python-dateutil" },
+ { name = "pytz" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ad/2f/44d1ae153a0e27be56be43465e5cb39b9650c781e001e7864389deb25090/croniter-6.0.0.tar.gz", hash = "sha256:37c504b313956114a983ece2c2b07790b1f1094fe9d81cc94739214748255577", size = 64481, upload-time = "2024-12-17T17:17:47.32Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/07/4b/290b4c3efd6417a8b0c284896de19b1d5855e6dbdb97d2a35e68fa42de85/croniter-6.0.0-py2.py3-none-any.whl", hash = "sha256:2f878c3856f17896979b2a4379ba1f09c83e374931ea15cc835c5dd2eee9b368", size = 25468, upload-time = "2024-12-17T17:17:45.359Z" },
+]
+
+[[package]]
+name = "cryptography"
+version = "46.0.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cffi", marker = "platform_python_implementation != 'PyPy'" },
+ { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/60/04/ee2a9e8542e4fa2773b81771ff8349ff19cdd56b7258a0cc442639052edb/cryptography-46.0.5.tar.gz", hash = "sha256:abace499247268e3757271b2f1e244b36b06f8515cf27c4d49468fc9eb16e93d", size = 750064, upload-time = "2026-02-10T19:18:38.255Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/f7/81/b0bb27f2ba931a65409c6b8a8b358a7f03c0e46eceacddff55f7c84b1f3b/cryptography-46.0.5-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:351695ada9ea9618b3500b490ad54c739860883df6c1f555e088eaf25b1bbaad", size = 7176289, upload-time = "2026-02-10T19:17:08.274Z" },
+ { url = "https://files.pythonhosted.org/packages/ff/9e/6b4397a3e3d15123de3b1806ef342522393d50736c13b20ec4c9ea6693a6/cryptography-46.0.5-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c18ff11e86df2e28854939acde2d003f7984f721eba450b56a200ad90eeb0e6b", size = 4275637, upload-time = "2026-02-10T19:17:10.53Z" },
+ { url = "https://files.pythonhosted.org/packages/63/e7/471ab61099a3920b0c77852ea3f0ea611c9702f651600397ac567848b897/cryptography-46.0.5-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d7e3d356b8cd4ea5aff04f129d5f66ebdc7b6f8eae802b93739ed520c47c79b", size = 4424742, upload-time = "2026-02-10T19:17:12.388Z" },
+ { url = "https://files.pythonhosted.org/packages/37/53/a18500f270342d66bf7e4d9f091114e31e5ee9e7375a5aba2e85a91e0044/cryptography-46.0.5-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:50bfb6925eff619c9c023b967d5b77a54e04256c4281b0e21336a130cd7fc263", size = 4277528, upload-time = "2026-02-10T19:17:13.853Z" },
+ { url = "https://files.pythonhosted.org/packages/22/29/c2e812ebc38c57b40e7c583895e73c8c5adb4d1e4a0cc4c5a4fdab2b1acc/cryptography-46.0.5-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:803812e111e75d1aa73690d2facc295eaefd4439be1023fefc4995eaea2af90d", size = 4947993, upload-time = "2026-02-10T19:17:15.618Z" },
+ { url = "https://files.pythonhosted.org/packages/6b/e7/237155ae19a9023de7e30ec64e5d99a9431a567407ac21170a046d22a5a3/cryptography-46.0.5-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3ee190460e2fbe447175cda91b88b84ae8322a104fc27766ad09428754a618ed", size = 4456855, upload-time = "2026-02-10T19:17:17.221Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/87/fc628a7ad85b81206738abbd213b07702bcbdada1dd43f72236ef3cffbb5/cryptography-46.0.5-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:f145bba11b878005c496e93e257c1e88f154d278d2638e6450d17e0f31e558d2", size = 3984635, upload-time = "2026-02-10T19:17:18.792Z" },
+ { url = "https://files.pythonhosted.org/packages/84/29/65b55622bde135aedf4565dc509d99b560ee4095e56989e815f8fd2aa910/cryptography-46.0.5-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:e9251e3be159d1020c4030bd2e5f84d6a43fe54b6c19c12f51cde9542a2817b2", size = 4277038, upload-time = "2026-02-10T19:17:20.256Z" },
+ { url = "https://files.pythonhosted.org/packages/bc/36/45e76c68d7311432741faf1fbf7fac8a196a0a735ca21f504c75d37e2558/cryptography-46.0.5-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:47fb8a66058b80e509c47118ef8a75d14c455e81ac369050f20ba0d23e77fee0", size = 4912181, upload-time = "2026-02-10T19:17:21.825Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/1a/c1ba8fead184d6e3d5afcf03d569acac5ad063f3ac9fb7258af158f7e378/cryptography-46.0.5-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:4c3341037c136030cb46e4b1e17b7418ea4cbd9dd207e4a6f3b2b24e0d4ac731", size = 4456482, upload-time = "2026-02-10T19:17:25.133Z" },
+ { url = "https://files.pythonhosted.org/packages/f9/e5/3fb22e37f66827ced3b902cf895e6a6bc1d095b5b26be26bd13c441fdf19/cryptography-46.0.5-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:890bcb4abd5a2d3f852196437129eb3667d62630333aacc13dfd470fad3aaa82", size = 4405497, upload-time = "2026-02-10T19:17:26.66Z" },
+ { url = "https://files.pythonhosted.org/packages/1a/df/9d58bb32b1121a8a2f27383fabae4d63080c7ca60b9b5c88be742be04ee7/cryptography-46.0.5-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:80a8d7bfdf38f87ca30a5391c0c9ce4ed2926918e017c29ddf643d0ed2778ea1", size = 4667819, upload-time = "2026-02-10T19:17:28.569Z" },
+ { url = "https://files.pythonhosted.org/packages/ea/ed/325d2a490c5e94038cdb0117da9397ece1f11201f425c4e9c57fe5b9f08b/cryptography-46.0.5-cp311-abi3-win32.whl", hash = "sha256:60ee7e19e95104d4c03871d7d7dfb3d22ef8a9b9c6778c94e1c8fcc8365afd48", size = 3028230, upload-time = "2026-02-10T19:17:30.518Z" },
+ { url = "https://files.pythonhosted.org/packages/e9/5a/ac0f49e48063ab4255d9e3b79f5def51697fce1a95ea1370f03dc9db76f6/cryptography-46.0.5-cp311-abi3-win_amd64.whl", hash = "sha256:38946c54b16c885c72c4f59846be9743d699eee2b69b6988e0a00a01f46a61a4", size = 3480909, upload-time = "2026-02-10T19:17:32.083Z" },
+ { url = "https://files.pythonhosted.org/packages/00/13/3d278bfa7a15a96b9dc22db5a12ad1e48a9eb3d40e1827ef66a5df75d0d0/cryptography-46.0.5-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:94a76daa32eb78d61339aff7952ea819b1734b46f73646a07decb40e5b3448e2", size = 7119287, upload-time = "2026-02-10T19:17:33.801Z" },
+ { url = "https://files.pythonhosted.org/packages/67/c8/581a6702e14f0898a0848105cbefd20c058099e2c2d22ef4e476dfec75d7/cryptography-46.0.5-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5be7bf2fb40769e05739dd0046e7b26f9d4670badc7b032d6ce4db64dddc0678", size = 4265728, upload-time = "2026-02-10T19:17:35.569Z" },
+ { url = "https://files.pythonhosted.org/packages/dd/4a/ba1a65ce8fc65435e5a849558379896c957870dd64fecea97b1ad5f46a37/cryptography-46.0.5-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fe346b143ff9685e40192a4960938545c699054ba11d4f9029f94751e3f71d87", size = 4408287, upload-time = "2026-02-10T19:17:36.938Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/67/8ffdbf7b65ed1ac224d1c2df3943553766914a8ca718747ee3871da6107e/cryptography-46.0.5-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:c69fd885df7d089548a42d5ec05be26050ebcd2283d89b3d30676eb32ff87dee", size = 4270291, upload-time = "2026-02-10T19:17:38.748Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/e5/f52377ee93bc2f2bba55a41a886fd208c15276ffbd2569f2ddc89d50e2c5/cryptography-46.0.5-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:8293f3dea7fc929ef7240796ba231413afa7b68ce38fd21da2995549f5961981", size = 4927539, upload-time = "2026-02-10T19:17:40.241Z" },
+ { url = "https://files.pythonhosted.org/packages/3b/02/cfe39181b02419bbbbcf3abdd16c1c5c8541f03ca8bda240debc467d5a12/cryptography-46.0.5-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1abfdb89b41c3be0365328a410baa9df3ff8a9110fb75e7b52e66803ddabc9a9", size = 4442199, upload-time = "2026-02-10T19:17:41.789Z" },
+ { url = "https://files.pythonhosted.org/packages/c0/96/2fcaeb4873e536cf71421a388a6c11b5bc846e986b2b069c79363dc1648e/cryptography-46.0.5-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:d66e421495fdb797610a08f43b05269e0a5ea7f5e652a89bfd5a7d3c1dee3648", size = 3960131, upload-time = "2026-02-10T19:17:43.379Z" },
+ { url = "https://files.pythonhosted.org/packages/d8/d2/b27631f401ddd644e94c5cf33c9a4069f72011821cf3dc7309546b0642a0/cryptography-46.0.5-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:4e817a8920bfbcff8940ecfd60f23d01836408242b30f1a708d93198393a80b4", size = 4270072, upload-time = "2026-02-10T19:17:45.481Z" },
+ { url = "https://files.pythonhosted.org/packages/f4/a7/60d32b0370dae0b4ebe55ffa10e8599a2a59935b5ece1b9f06edb73abdeb/cryptography-46.0.5-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:68f68d13f2e1cb95163fa3b4db4bf9a159a418f5f6e7242564fc75fcae667fd0", size = 4892170, upload-time = "2026-02-10T19:17:46.997Z" },
+ { url = "https://files.pythonhosted.org/packages/d2/b9/cf73ddf8ef1164330eb0b199a589103c363afa0cf794218c24d524a58eab/cryptography-46.0.5-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:a3d1fae9863299076f05cb8a778c467578262fae09f9dc0ee9b12eb4268ce663", size = 4441741, upload-time = "2026-02-10T19:17:48.661Z" },
+ { url = "https://files.pythonhosted.org/packages/5f/eb/eee00b28c84c726fe8fa0158c65afe312d9c3b78d9d01daf700f1f6e37ff/cryptography-46.0.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c4143987a42a2397f2fc3b4d7e3a7d313fbe684f67ff443999e803dd75a76826", size = 4396728, upload-time = "2026-02-10T19:17:50.058Z" },
+ { url = "https://files.pythonhosted.org/packages/65/f4/6bc1a9ed5aef7145045114b75b77c2a8261b4d38717bd8dea111a63c3442/cryptography-46.0.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:7d731d4b107030987fd61a7f8ab512b25b53cef8f233a97379ede116f30eb67d", size = 4652001, upload-time = "2026-02-10T19:17:51.54Z" },
+ { url = "https://files.pythonhosted.org/packages/86/ef/5d00ef966ddd71ac2e6951d278884a84a40ffbd88948ef0e294b214ae9e4/cryptography-46.0.5-cp314-cp314t-win32.whl", hash = "sha256:c3bcce8521d785d510b2aad26ae2c966092b7daa8f45dd8f44734a104dc0bc1a", size = 3003637, upload-time = "2026-02-10T19:17:52.997Z" },
+ { url = "https://files.pythonhosted.org/packages/b7/57/f3f4160123da6d098db78350fdfd9705057aad21de7388eacb2401dceab9/cryptography-46.0.5-cp314-cp314t-win_amd64.whl", hash = "sha256:4d8ae8659ab18c65ced284993c2265910f6c9e650189d4e3f68445ef82a810e4", size = 3469487, upload-time = "2026-02-10T19:17:54.549Z" },
+ { url = "https://files.pythonhosted.org/packages/e2/fa/a66aa722105ad6a458bebd64086ca2b72cdd361fed31763d20390f6f1389/cryptography-46.0.5-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:4108d4c09fbbf2789d0c926eb4152ae1760d5a2d97612b92d508d96c861e4d31", size = 7170514, upload-time = "2026-02-10T19:17:56.267Z" },
+ { url = "https://files.pythonhosted.org/packages/0f/04/c85bdeab78c8bc77b701bf0d9bdcf514c044e18a46dcff330df5448631b0/cryptography-46.0.5-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7d1f30a86d2757199cb2d56e48cce14deddf1f9c95f1ef1b64ee91ea43fe2e18", size = 4275349, upload-time = "2026-02-10T19:17:58.419Z" },
+ { url = "https://files.pythonhosted.org/packages/5c/32/9b87132a2f91ee7f5223b091dc963055503e9b442c98fc0b8a5ca765fab0/cryptography-46.0.5-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:039917b0dc418bb9f6edce8a906572d69e74bd330b0b3fea4f79dab7f8ddd235", size = 4420667, upload-time = "2026-02-10T19:18:00.619Z" },
+ { url = "https://files.pythonhosted.org/packages/a1/a6/a7cb7010bec4b7c5692ca6f024150371b295ee1c108bdc1c400e4c44562b/cryptography-46.0.5-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ba2a27ff02f48193fc4daeadf8ad2590516fa3d0adeeb34336b96f7fa64c1e3a", size = 4276980, upload-time = "2026-02-10T19:18:02.379Z" },
+ { url = "https://files.pythonhosted.org/packages/8e/7c/c4f45e0eeff9b91e3f12dbd0e165fcf2a38847288fcfd889deea99fb7b6d/cryptography-46.0.5-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:61aa400dce22cb001a98014f647dc21cda08f7915ceb95df0c9eaf84b4b6af76", size = 4939143, upload-time = "2026-02-10T19:18:03.964Z" },
+ { url = "https://files.pythonhosted.org/packages/37/19/e1b8f964a834eddb44fa1b9a9976f4e414cbb7aa62809b6760c8803d22d1/cryptography-46.0.5-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3ce58ba46e1bc2aac4f7d9290223cead56743fa6ab94a5d53292ffaac6a91614", size = 4453674, upload-time = "2026-02-10T19:18:05.588Z" },
+ { url = "https://files.pythonhosted.org/packages/db/ed/db15d3956f65264ca204625597c410d420e26530c4e2943e05a0d2f24d51/cryptography-46.0.5-cp38-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:420d0e909050490d04359e7fdb5ed7e667ca5c3c402b809ae2563d7e66a92229", size = 3978801, upload-time = "2026-02-10T19:18:07.167Z" },
+ { url = "https://files.pythonhosted.org/packages/41/e2/df40a31d82df0a70a0daf69791f91dbb70e47644c58581d654879b382d11/cryptography-46.0.5-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:582f5fcd2afa31622f317f80426a027f30dc792e9c80ffee87b993200ea115f1", size = 4276755, upload-time = "2026-02-10T19:18:09.813Z" },
+ { url = "https://files.pythonhosted.org/packages/33/45/726809d1176959f4a896b86907b98ff4391a8aa29c0aaaf9450a8a10630e/cryptography-46.0.5-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:bfd56bb4b37ed4f330b82402f6f435845a5f5648edf1ad497da51a8452d5d62d", size = 4901539, upload-time = "2026-02-10T19:18:11.263Z" },
+ { url = "https://files.pythonhosted.org/packages/99/0f/a3076874e9c88ecb2ecc31382f6e7c21b428ede6f55aafa1aa272613e3cd/cryptography-46.0.5-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:a3d507bb6a513ca96ba84443226af944b0f7f47dcc9a399d110cd6146481d24c", size = 4452794, upload-time = "2026-02-10T19:18:12.914Z" },
+ { url = "https://files.pythonhosted.org/packages/02/ef/ffeb542d3683d24194a38f66ca17c0a4b8bf10631feef44a7ef64e631b1a/cryptography-46.0.5-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9f16fbdf4da055efb21c22d81b89f155f02ba420558db21288b3d0035bafd5f4", size = 4404160, upload-time = "2026-02-10T19:18:14.375Z" },
+ { url = "https://files.pythonhosted.org/packages/96/93/682d2b43c1d5f1406ed048f377c0fc9fc8f7b0447a478d5c65ab3d3a66eb/cryptography-46.0.5-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:ced80795227d70549a411a4ab66e8ce307899fad2220ce5ab2f296e687eacde9", size = 4667123, upload-time = "2026-02-10T19:18:15.886Z" },
+ { url = "https://files.pythonhosted.org/packages/45/2d/9c5f2926cb5300a8eefc3f4f0b3f3df39db7f7ce40c8365444c49363cbda/cryptography-46.0.5-cp38-abi3-win32.whl", hash = "sha256:02f547fce831f5096c9a567fd41bc12ca8f11df260959ecc7c3202555cc47a72", size = 3010220, upload-time = "2026-02-10T19:18:17.361Z" },
+ { url = "https://files.pythonhosted.org/packages/48/ef/0c2f4a8e31018a986949d34a01115dd057bf536905dca38897bacd21fac3/cryptography-46.0.5-cp38-abi3-win_amd64.whl", hash = "sha256:556e106ee01aa13484ce9b0239bca667be5004efb0aabbed28d353df86445595", size = 3467050, upload-time = "2026-02-10T19:18:18.899Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/dd/2d9fdb07cebdf3d51179730afb7d5e576153c6744c3ff8fded23030c204e/cryptography-46.0.5-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:3b4995dc971c9fb83c25aa44cf45f02ba86f71ee600d81091c2f0cbae116b06c", size = 3476964, upload-time = "2026-02-10T19:18:20.687Z" },
+ { url = "https://files.pythonhosted.org/packages/e9/6f/6cc6cc9955caa6eaf83660b0da2b077c7fe8ff9950a3c5e45d605038d439/cryptography-46.0.5-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:bc84e875994c3b445871ea7181d424588171efec3e185dced958dad9e001950a", size = 4218321, upload-time = "2026-02-10T19:18:22.349Z" },
+ { url = "https://files.pythonhosted.org/packages/3e/5d/c4da701939eeee699566a6c1367427ab91a8b7088cc2328c09dbee940415/cryptography-46.0.5-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:2ae6971afd6246710480e3f15824ed3029a60fc16991db250034efd0b9fb4356", size = 4381786, upload-time = "2026-02-10T19:18:24.529Z" },
+ { url = "https://files.pythonhosted.org/packages/ac/97/a538654732974a94ff96c1db621fa464f455c02d4bb7d2652f4edc21d600/cryptography-46.0.5-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:d861ee9e76ace6cf36a6a89b959ec08e7bc2493ee39d07ffe5acb23ef46d27da", size = 4217990, upload-time = "2026-02-10T19:18:25.957Z" },
+ { url = "https://files.pythonhosted.org/packages/ae/11/7e500d2dd3ba891197b9efd2da5454b74336d64a7cc419aa7327ab74e5f6/cryptography-46.0.5-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:2b7a67c9cd56372f3249b39699f2ad479f6991e62ea15800973b956f4b73e257", size = 4381252, upload-time = "2026-02-10T19:18:27.496Z" },
+ { url = "https://files.pythonhosted.org/packages/bc/58/6b3d24e6b9bc474a2dcdee65dfd1f008867015408a271562e4b690561a4d/cryptography-46.0.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:8456928655f856c6e1533ff59d5be76578a7157224dbd9ce6872f25055ab9ab7", size = 3407605, upload-time = "2026-02-10T19:18:29.233Z" },
+]
+
+[[package]]
+name = "cyclopts"
+version = "4.5.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "attrs" },
+ { name = "docstring-parser" },
+ { name = "rich" },
+ { name = "rich-rst" },
+ { name = "tomli", marker = "python_full_version < '3.11'" },
+ { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a5/16/06e35c217334930ff7c476ce1c8e74ed786fa3ef6742e59a1458e2412290/cyclopts-4.5.3.tar.gz", hash = "sha256:35fa70971204c450d9668646a6ca372eb5fa3070fbe8dd51c5b4b31e65198f2d", size = 162437, upload-time = "2026-02-16T15:07:11.96Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/3a/1f/d8bce383a90d8a6a11033327777afa4d4d611ec11869284adb6f48152906/cyclopts-4.5.3-py3-none-any.whl", hash = "sha256:50af3085bb15d4a6f2582dd383dad5e4ba6a0d4d4c64ee63326d881a752a6919", size = 200231, upload-time = "2026-02-16T15:07:13.045Z" },
+]
+
+[[package]]
+name = "diskcache"
+version = "5.6.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3f/21/1c1ffc1a039ddcc459db43cc108658f32c57d271d7289a2794e401d0fdb6/diskcache-5.6.3.tar.gz", hash = "sha256:2c3a3fa2743d8535d832ec61c2054a1641f41775aa7c556758a109941e33e4fc", size = 67916, upload-time = "2023-08-31T06:12:00.316Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" },
+]
+
[[package]]
name = "distro"
version = "1.9.0"
@@ -379,6 +564,46 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" },
]
+[[package]]
+name = "dnspython"
+version = "2.8.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" },
+]
+
+[[package]]
+name = "docstring-parser"
+version = "0.17.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b2/9d/c3b43da9515bd270df0f80548d9944e389870713cc1fe2b8fb35fe2bcefd/docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912", size = 27442, upload-time = "2025-07-21T07:35:01.868Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" },
+]
+
+[[package]]
+name = "docutils"
+version = "0.22.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ae/b6/03bb70946330e88ffec97aefd3ea75ba575cb2e762061e0e62a213befee8/docutils-0.22.4.tar.gz", hash = "sha256:4db53b1fde9abecbb74d91230d32ab626d94f6badfc575d6db9194a49df29968", size = 2291750, upload-time = "2025-12-18T19:00:26.443Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" },
+]
+
+[[package]]
+name = "email-validator"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "dnspython" },
+ { name = "idna" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" },
+]
+
[[package]]
name = "exceptiongroup"
version = "1.3.1"
@@ -391,6 +616,25 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" },
]
+[[package]]
+name = "fakeredis"
+version = "2.34.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "redis" },
+ { name = "sortedcontainers" },
+ { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d8/44/c403963727d707e03f49a417712b0a23e853d33ae50729679040b6cfe281/fakeredis-2.34.0.tar.gz", hash = "sha256:72bc51a7ab39bedf5004f0cf1b5206822619c1be8c2657fd878d1f4250256c57", size = 177156, upload-time = "2026-02-16T15:56:34.318Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/1a/8e/af19c00753c432355f9b76cec3ab0842578de43ba575e82735b18c1b3ec9/fakeredis-2.34.0-py3-none-any.whl", hash = "sha256:bc45d362c6cc3a537f8287372d8ea532538dfbe7f5d635d0905d7b3464ec51d2", size = 122063, upload-time = "2026-02-16T15:56:21.227Z" },
+]
+
+[package.optional-dependencies]
+lua = [
+ { name = "lupa" },
+]
+
[[package]]
name = "fastapi"
version = "0.124.4"
@@ -406,6 +650,35 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3e/57/aa70121b5008f44031be645a61a7c4abc24e0e888ad3fc8fda916f4d188e/fastapi-0.124.4-py3-none-any.whl", hash = "sha256:6d1e703698443ccb89e50abe4893f3c84d9d6689c0cf1ca4fad6d3c15cf69f15", size = 113281, upload-time = "2025-12-12T15:00:42.44Z" },
]
+[[package]]
+name = "fastmcp"
+version = "2.14.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "authlib" },
+ { name = "cyclopts" },
+ { name = "exceptiongroup" },
+ { name = "httpx" },
+ { name = "jsonref" },
+ { name = "jsonschema-path" },
+ { name = "mcp" },
+ { name = "openapi-pydantic" },
+ { name = "packaging" },
+ { name = "platformdirs" },
+ { name = "py-key-value-aio", extra = ["disk", "keyring", "memory"] },
+ { name = "pydantic", extra = ["email"] },
+ { name = "pydocket" },
+ { name = "pyperclip" },
+ { name = "python-dotenv" },
+ { name = "rich" },
+ { name = "uvicorn" },
+ { name = "websockets" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3b/32/982678d44f13849530a74ab101ed80e060c2ee6cf87471f062dcf61705fd/fastmcp-2.14.5.tar.gz", hash = "sha256:38944dc582c541d55357082bda2241cedb42cd3a78faea8a9d6a2662c62a42d7", size = 8296329, upload-time = "2026-02-03T15:35:21.005Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/e5/c1/1a35ec68ff76ea8443aa115b18bcdee748a4ada2124537ee90522899ff9f/fastmcp-2.14.5-py3-none-any.whl", hash = "sha256:d81e8ec813f5089d3624bec93944beaefa86c0c3a4ef1111cbef676a761ebccf", size = 417784, upload-time = "2026-02-03T15:35:18.489Z" },
+]
+
[[package]]
name = "fastuuid"
version = "0.14.0"
@@ -735,6 +1008,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
]
+[[package]]
+name = "httpx-sse"
+version = "0.4.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" },
+]
+
[[package]]
name = "huggingface-hub"
version = "1.3.1"
@@ -787,12 +1069,48 @@ wheels = [
]
[[package]]
-name = "isort"
+name = "jaraco-classes"
+version = "3.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "more-itertools" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" },
+]
+
+[[package]]
+name = "jaraco-context"
version = "6.1.0"
source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1e/82/fa43935523efdfcce6abbae9da7f372b627b27142c3419fcf13bf5b0c397/isort-6.1.0.tar.gz", hash = "sha256:9b8f96a14cfee0677e78e941ff62f03769a06d412aabb9e2a90487b3b7e8d481", size = 824325, upload-time = "2025-10-01T16:26:45.027Z" }
+dependencies = [
+ { name = "backports-tarfile", marker = "python_full_version < '3.12'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/cb/9c/a788f5bb29c61e456b8ee52ce76dbdd32fd72cd73dd67bc95f42c7a8d13c/jaraco_context-6.1.0.tar.gz", hash = "sha256:129a341b0a85a7db7879e22acd66902fda67882db771754574338898b2d5d86f", size = 15850, upload-time = "2026-01-13T02:53:53.847Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/7f/cc/9b681a170efab4868a032631dea1e8446d8ec718a7f657b94d49d1a12643/isort-6.1.0-py3-none-any.whl", hash = "sha256:58d8927ecce74e5087aef019f778d4081a3b6c98f15a80ba35782ca8a2097784", size = 94329, upload-time = "2025-10-01T16:26:43.291Z" },
+ { url = "https://files.pythonhosted.org/packages/8d/48/aa685dbf1024c7bd82bede569e3a85f82c32fd3d79ba5fea578f0159571a/jaraco_context-6.1.0-py3-none-any.whl", hash = "sha256:a43b5ed85815223d0d3cfdb6d7ca0d2bc8946f28f30b6f3216bda070f68badda", size = 7065, upload-time = "2026-01-13T02:53:53.031Z" },
+]
+
+[[package]]
+name = "jaraco-functools"
+version = "4.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "more-itertools" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/0f/27/056e0638a86749374d6f57d0b0db39f29509cce9313cf91bdc0ac4d91084/jaraco_functools-4.4.0.tar.gz", hash = "sha256:da21933b0417b89515562656547a77b4931f98176eb173644c0d35032a33d6bb", size = 19943, upload-time = "2025-12-21T09:29:43.6Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/fd/c4/813bb09f0985cb21e959f21f2464169eca882656849adf727ac7bb7e1767/jaraco_functools-4.4.0-py3-none-any.whl", hash = "sha256:9eec1e36f45c818d9bf307c8948eb03b2b56cd44087b3cdc989abca1f20b9176", size = 10481, upload-time = "2025-12-21T09:29:42.27Z" },
+]
+
+[[package]]
+name = "jeepney"
+version = "0.9.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" },
]
[[package]]
@@ -904,6 +1222,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/2f/9c/6753e6522b8d0ef07d3a3d239426669e984fb0eba15a315cdbc1253904e4/jiter-0.12.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c24e864cb30ab82311c6425655b0cdab0a98c5d973b065c66a3f020740c2324c", size = 346110, upload-time = "2025-11-09T20:49:21.817Z" },
]
+[[package]]
+name = "jsonref"
+version = "1.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/aa/0d/c1f3277e90ccdb50d33ed5ba1ec5b3f0a242ed8c1b1a85d3afeb68464dca/jsonref-1.1.0.tar.gz", hash = "sha256:32fe8e1d85af0fdefbebce950af85590b22b60f9e95443176adbde4e1ecea552", size = 8814, upload-time = "2023-01-16T16:10:04.455Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/0c/ec/e1db9922bceb168197a558a2b8c03a7963f1afe93517ddd3cf99f202f996/jsonref-1.1.0-py3-none-any.whl", hash = "sha256:590dc7773df6c21cbf948b5dac07a72a251db28b0238ceecce0a2abfa8ec30a9", size = 9425, upload-time = "2023-01-16T16:10:02.255Z" },
+]
+
[[package]]
name = "jsonschema"
version = "4.26.0"
@@ -919,6 +1246,21 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" },
]
+[[package]]
+name = "jsonschema-path"
+version = "0.3.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "pathable" },
+ { name = "pyyaml" },
+ { name = "referencing" },
+ { name = "requests" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6e/45/41ebc679c2a4fced6a722f624c18d658dee42612b83ea24c1caf7c0eb3a8/jsonschema_path-0.3.4.tar.gz", hash = "sha256:8365356039f16cc65fddffafda5f58766e34bebab7d6d105616ab52bc4297001", size = 11159, upload-time = "2025-01-24T14:33:16.547Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/cb/58/3485da8cb93d2f393bce453adeef16896751f14ba3e2024bc21dc9597646/jsonschema_path-0.3.4-py3-none-any.whl", hash = "sha256:f502191fdc2b22050f9a81c9237be9d27145b9001c55842bece5e94e382e52f8", size = 14810, upload-time = "2025-01-24T14:33:14.652Z" },
+]
+
[[package]]
name = "jsonschema-specifications"
version = "2025.9.1"
@@ -931,6 +1273,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" },
]
+[[package]]
+name = "keyring"
+version = "25.7.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "importlib-metadata", marker = "python_full_version < '3.12'" },
+ { name = "jaraco-classes" },
+ { name = "jaraco-context" },
+ { name = "jaraco-functools" },
+ { name = "jeepney", marker = "sys_platform == 'linux'" },
+ { name = "pywin32-ctypes", marker = "sys_platform == 'win32'" },
+ { name = "secretstorage", marker = "sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" },
+]
+
[[package]]
name = "litellm"
version = "1.80.15"
@@ -955,6 +1315,80 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/df/3b/b1bd693721ccb3c9a37c8233d019a643ac57bef5a93f279e5a63839ee4db/litellm-1.80.15-py3-none-any.whl", hash = "sha256:f354e49456985a235b9ed99df1c19d686d30501f96e68882dcc5b29b1e7c59d9", size = 11670707, upload-time = "2026-01-11T18:31:41.67Z" },
]
+[[package]]
+name = "lupa"
+version = "2.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b8/1c/191c3e6ec6502e3dbe25a53e27f69a5daeac3e56de1f73c0138224171ead/lupa-2.6.tar.gz", hash = "sha256:9a770a6e89576be3447668d7ced312cd6fd41d3c13c2462c9dc2c2ab570e45d9", size = 7240282, upload-time = "2025-10-24T07:20:29.738Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/a1/15/713cab5d0dfa4858f83b99b3e0329072df33dc14fc3ebbaa017e0f9755c4/lupa-2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6b3dabda836317e63c5ad052826e156610f356a04b3003dfa0dbe66b5d54d671", size = 954828, upload-time = "2025-10-24T07:17:15.726Z" },
+ { url = "https://files.pythonhosted.org/packages/2e/71/704740cbc6e587dd6cc8dabf2f04820ac6a671784e57cc3c29db795476db/lupa-2.6-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:8726d1c123bbe9fbb974ce29825e94121824e66003038ff4532c14cc2ed0c51c", size = 1919259, upload-time = "2025-10-24T07:17:18.586Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/18/f248341c423c5d48837e35584c6c3eb4acab7e722b6057d7b3e28e42dae8/lupa-2.6-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:f4e159e7d814171199b246f9235ca8961f6461ea8c1165ab428afa13c9289a94", size = 984998, upload-time = "2025-10-24T07:17:20.428Z" },
+ { url = "https://files.pythonhosted.org/packages/44/1e/8a4bd471e018aad76bcb9455d298c2c96d82eced20f2ae8fcec8cd800948/lupa-2.6-cp310-cp310-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:202160e80dbfddfb79316692a563d843b767e0f6787bbd1c455f9d54052efa6c", size = 1174871, upload-time = "2025-10-24T07:17:22.755Z" },
+ { url = "https://files.pythonhosted.org/packages/2a/5c/3a3f23fd6a91b0986eea1ceaf82ad3f9b958fe3515a9981fb9c4eb046c8b/lupa-2.6-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5deede7c5b36ab64f869dae4831720428b67955b0bb186c8349cf6ea121c852b", size = 1057471, upload-time = "2025-10-24T07:17:24.908Z" },
+ { url = "https://files.pythonhosted.org/packages/45/ac/01be1fed778fb0c8f46ee8cbe344e4d782f6806fac12717f08af87aa4355/lupa-2.6-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:86f04901f920bbf7c0cac56807dc9597e42347123e6f1f3ca920f15f54188ce5", size = 2100592, upload-time = "2025-10-24T07:17:27.089Z" },
+ { url = "https://files.pythonhosted.org/packages/3f/6c/1a05bb873e30830f8574e10cd0b4cdbc72e9dbad2a09e25810b5e3b1f75d/lupa-2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:6deef8f851d6afb965c84849aa5b8c38856942df54597a811ce0369ced678610", size = 1081396, upload-time = "2025-10-24T07:17:29.064Z" },
+ { url = "https://files.pythonhosted.org/packages/a2/c2/a19dd80d6dc98b39bbf8135b8198e38aa7ca3360b720eac68d1d7e9286b5/lupa-2.6-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:21f2b5549681c2a13b1170a26159d30875d367d28f0247b81ca347222c755038", size = 1192007, upload-time = "2025-10-24T07:17:31.362Z" },
+ { url = "https://files.pythonhosted.org/packages/4f/43/e1b297225c827f55752e46fdbfb021c8982081b0f24490e42776ea69ae3b/lupa-2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:66eea57630eab5e6f49fdc5d7811c0a2a41f2011be4ea56a087ea76112011eb7", size = 2196661, upload-time = "2025-10-24T07:17:33.484Z" },
+ { url = "https://files.pythonhosted.org/packages/2e/8f/2272d429a7fa9dc8dbd6e9c5c9073a03af6007eb22a4c78829fec6a34b80/lupa-2.6-cp310-cp310-win32.whl", hash = "sha256:60a403de8cab262a4fe813085dd77010effa6e2eb1886db2181df803140533b1", size = 1412738, upload-time = "2025-10-24T07:17:35.11Z" },
+ { url = "https://files.pythonhosted.org/packages/35/2a/1708911271dd49ad87b4b373b5a4b0e0a0516d3d2af7b76355946c7ee171/lupa-2.6-cp310-cp310-win_amd64.whl", hash = "sha256:e4656a39d93dfa947cf3db56dc16c7916cb0cc8024acd3a952071263f675df64", size = 1656898, upload-time = "2025-10-24T07:17:36.949Z" },
+ { url = "https://files.pythonhosted.org/packages/ca/29/1f66907c1ebf1881735afa695e646762c674f00738ebf66d795d59fc0665/lupa-2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6d988c0f9331b9f2a5a55186701a25444ab10a1432a1021ee58011499ecbbdd5", size = 962875, upload-time = "2025-10-24T07:17:39.107Z" },
+ { url = "https://files.pythonhosted.org/packages/e6/67/4a748604be360eb9c1c215f6a0da921cd1a2b44b2c5951aae6fb83019d3a/lupa-2.6-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:ebe1bbf48259382c72a6fe363dea61a0fd6fe19eab95e2ae881e20f3654587bf", size = 1935390, upload-time = "2025-10-24T07:17:41.427Z" },
+ { url = "https://files.pythonhosted.org/packages/ac/0c/8ef9ee933a350428b7bdb8335a37ef170ab0bb008bbf9ca8f4f4310116b6/lupa-2.6-cp311-cp311-macosx_11_0_x86_64.whl", hash = "sha256:a8fcee258487cf77cdd41560046843bb38c2e18989cd19671dd1e2596f798306", size = 992193, upload-time = "2025-10-24T07:17:43.231Z" },
+ { url = "https://files.pythonhosted.org/packages/65/46/e6c7facebdb438db8a65ed247e56908818389c1a5abbf6a36aab14f1057d/lupa-2.6-cp311-cp311-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:561a8e3be800827884e767a694727ed8482d066e0d6edfcbf423b05e63b05535", size = 1165844, upload-time = "2025-10-24T07:17:45.437Z" },
+ { url = "https://files.pythonhosted.org/packages/1c/26/9f1154c6c95f175ccbf96aa96c8f569c87f64f463b32473e839137601a8b/lupa-2.6-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:af880a62d47991cae78b8e9905c008cbfdc4a3a9723a66310c2634fc7644578c", size = 1048069, upload-time = "2025-10-24T07:17:47.181Z" },
+ { url = "https://files.pythonhosted.org/packages/68/67/2cc52ab73d6af81612b2ea24c870d3fa398443af8e2875e5befe142398b1/lupa-2.6-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:80b22923aa4023c86c0097b235615f89d469a0c4eee0489699c494d3367c4c85", size = 2079079, upload-time = "2025-10-24T07:17:49.755Z" },
+ { url = "https://files.pythonhosted.org/packages/2e/dc/f843f09bbf325f6e5ee61730cf6c3409fc78c010d968c7c78acba3019ca7/lupa-2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:153d2cc6b643f7efb9cfc0c6bb55ec784d5bac1a3660cfc5b958a7b8f38f4a75", size = 1071428, upload-time = "2025-10-24T07:17:51.991Z" },
+ { url = "https://files.pythonhosted.org/packages/2e/60/37533a8d85bf004697449acb97ecdacea851acad28f2ad3803662487dd2a/lupa-2.6-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:3fa8777e16f3ded50b72967dc17e23f5a08e4f1e2c9456aff2ebdb57f5b2869f", size = 1181756, upload-time = "2025-10-24T07:17:53.752Z" },
+ { url = "https://files.pythonhosted.org/packages/e4/f2/cf29b20dbb4927b6a3d27c339ac5d73e74306ecc28c8e2c900b2794142ba/lupa-2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:8dbdcbe818c02a2f56f5ab5ce2de374dab03e84b25266cfbaef237829bc09b3f", size = 2175687, upload-time = "2025-10-24T07:17:56.228Z" },
+ { url = "https://files.pythonhosted.org/packages/94/7c/050e02f80c7131b63db1474bff511e63c545b5a8636a24cbef3fc4da20b6/lupa-2.6-cp311-cp311-win32.whl", hash = "sha256:defaf188fde8f7a1e5ce3a5e6d945e533b8b8d547c11e43b96c9b7fe527f56dc", size = 1412592, upload-time = "2025-10-24T07:17:59.062Z" },
+ { url = "https://files.pythonhosted.org/packages/6f/9a/6f2af98aa5d771cea661f66c8eb8f53772ec1ab1dfbce24126cfcd189436/lupa-2.6-cp311-cp311-win_amd64.whl", hash = "sha256:9505ae600b5c14f3e17e70f87f88d333717f60411faca1ddc6f3e61dce85fa9e", size = 1669194, upload-time = "2025-10-24T07:18:01.647Z" },
+ { url = "https://files.pythonhosted.org/packages/94/86/ce243390535c39d53ea17ccf0240815e6e457e413e40428a658ea4ee4b8d/lupa-2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:47ce718817ef1cc0c40d87c3d5ae56a800d61af00fbc0fad1ca9be12df2f3b56", size = 951707, upload-time = "2025-10-24T07:18:03.884Z" },
+ { url = "https://files.pythonhosted.org/packages/86/85/cedea5e6cbeb54396fdcc55f6b741696f3f036d23cfaf986d50d680446da/lupa-2.6-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:7aba985b15b101495aa4b07112cdc08baa0c545390d560ad5cfde2e9e34f4d58", size = 1916703, upload-time = "2025-10-24T07:18:05.6Z" },
+ { url = "https://files.pythonhosted.org/packages/24/be/3d6b5f9a8588c01a4d88129284c726017b2089f3a3fd3ba8bd977292fea0/lupa-2.6-cp312-cp312-macosx_11_0_x86_64.whl", hash = "sha256:b766f62f95b2739f2248977d29b0722e589dcf4f0ccfa827ccbd29f0148bd2e5", size = 985152, upload-time = "2025-10-24T07:18:08.561Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/23/9f9a05beee5d5dce9deca4cb07c91c40a90541fc0a8e09db4ee670da550f/lupa-2.6-cp312-cp312-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:00a934c23331f94cb51760097ebfab14b005d55a6b30a2b480e3c53dd2fa290d", size = 1159599, upload-time = "2025-10-24T07:18:10.346Z" },
+ { url = "https://files.pythonhosted.org/packages/40/4e/e7c0583083db9d7f1fd023800a9767d8e4391e8330d56c2373d890ac971b/lupa-2.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21de9f38bd475303e34a042b7081aabdf50bd9bafd36ce4faea2f90fd9f15c31", size = 1038686, upload-time = "2025-10-24T07:18:12.112Z" },
+ { url = "https://files.pythonhosted.org/packages/1c/9f/5a4f7d959d4feba5e203ff0c31889e74d1ca3153122be4a46dca7d92bf7c/lupa-2.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cf3bda96d3fc41237e964a69c23647d50d4e28421111360274d4799832c560e9", size = 2071956, upload-time = "2025-10-24T07:18:14.572Z" },
+ { url = "https://files.pythonhosted.org/packages/92/34/2f4f13ca65d01169b1720176aedc4af17bc19ee834598c7292db232cb6dc/lupa-2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5a76ead245da54801a81053794aa3975f213221f6542d14ec4b859ee2e7e0323", size = 1057199, upload-time = "2025-10-24T07:18:16.379Z" },
+ { url = "https://files.pythonhosted.org/packages/35/2a/5f7d2eebec6993b0dcd428e0184ad71afb06a45ba13e717f6501bfed1da3/lupa-2.6-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8dd0861741caa20886ddbda0a121d8e52fb9b5bb153d82fa9bba796962bf30e8", size = 1173693, upload-time = "2025-10-24T07:18:18.153Z" },
+ { url = "https://files.pythonhosted.org/packages/e4/29/089b4d2f8e34417349af3904bb40bec40b65c8731f45e3fd8d497ca573e5/lupa-2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:239e63948b0b23023f81d9a19a395e768ed3da6a299f84e7963b8f813f6e3f9c", size = 2164394, upload-time = "2025-10-24T07:18:20.403Z" },
+ { url = "https://files.pythonhosted.org/packages/f3/1b/79c17b23c921f81468a111cad843b076a17ef4b684c4a8dff32a7969c3f0/lupa-2.6-cp312-cp312-win32.whl", hash = "sha256:325894e1099499e7a6f9c351147661a2011887603c71086d36fe0f964d52d1ce", size = 1420647, upload-time = "2025-10-24T07:18:23.368Z" },
+ { url = "https://files.pythonhosted.org/packages/b8/15/5121e68aad3584e26e1425a5c9a79cd898f8a152292059e128c206ee817c/lupa-2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c735a1ce8ee60edb0fe71d665f1e6b7c55c6021f1d340eb8c865952c602cd36f", size = 1688529, upload-time = "2025-10-24T07:18:25.523Z" },
+ { url = "https://files.pythonhosted.org/packages/28/1d/21176b682ca5469001199d8b95fa1737e29957a3d185186e7a8b55345f2e/lupa-2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:663a6e58a0f60e7d212017d6678639ac8df0119bc13c2145029dcba084391310", size = 947232, upload-time = "2025-10-24T07:18:27.878Z" },
+ { url = "https://files.pythonhosted.org/packages/ce/4c/d327befb684660ca13cf79cd1f1d604331808f9f1b6fb6bf57832f8edf80/lupa-2.6-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:d1f5afda5c20b1f3217a80e9bc1b77037f8a6eb11612fd3ada19065303c8f380", size = 1908625, upload-time = "2025-10-24T07:18:29.944Z" },
+ { url = "https://files.pythonhosted.org/packages/66/8e/ad22b0a19454dfd08662237a84c792d6d420d36b061f239e084f29d1a4f3/lupa-2.6-cp313-cp313-macosx_11_0_x86_64.whl", hash = "sha256:26f2b3c085fe76e9119e48c1013c1cccdc1f51585d456858290475aa38e7089e", size = 981057, upload-time = "2025-10-24T07:18:31.553Z" },
+ { url = "https://files.pythonhosted.org/packages/5c/48/74859073ab276bd0566c719f9ca0108b0cfc1956ca0d68678d117d47d155/lupa-2.6-cp313-cp313-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:60d2f902c7b96fb8ab98493dcff315e7bb4d0b44dc9dd76eb37de575025d5685", size = 1156227, upload-time = "2025-10-24T07:18:33.981Z" },
+ { url = "https://files.pythonhosted.org/packages/09/6c/0e9ded061916877253c2266074060eb71ed99fb21d73c8c114a76725bce2/lupa-2.6-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a02d25dee3a3250967c36590128d9220ae02f2eda166a24279da0b481519cbff", size = 1035752, upload-time = "2025-10-24T07:18:36.32Z" },
+ { url = "https://files.pythonhosted.org/packages/dd/ef/f8c32e454ef9f3fe909f6c7d57a39f950996c37a3deb7b391fec7903dab7/lupa-2.6-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6eae1ee16b886b8914ff292dbefbf2f48abfbdee94b33a88d1d5475e02423203", size = 2069009, upload-time = "2025-10-24T07:18:38.072Z" },
+ { url = "https://files.pythonhosted.org/packages/53/dc/15b80c226a5225815a890ee1c11f07968e0aba7a852df41e8ae6fe285063/lupa-2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b0edd5073a4ee74ab36f74fe61450148e6044f3952b8d21248581f3c5d1a58be", size = 1056301, upload-time = "2025-10-24T07:18:40.165Z" },
+ { url = "https://files.pythonhosted.org/packages/31/14/2086c1425c985acfb30997a67e90c39457122df41324d3c179d6ee2292c6/lupa-2.6-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:0c53ee9f22a8a17e7d4266ad48e86f43771951797042dd51d1494aaa4f5f3f0a", size = 1170673, upload-time = "2025-10-24T07:18:42.426Z" },
+ { url = "https://files.pythonhosted.org/packages/10/e5/b216c054cf86576c0191bf9a9f05de6f7e8e07164897d95eea0078dca9b2/lupa-2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:de7c0f157a9064a400d828789191a96da7f4ce889969a588b87ec80de9b14772", size = 2162227, upload-time = "2025-10-24T07:18:46.112Z" },
+ { url = "https://files.pythonhosted.org/packages/59/2f/33ecb5bedf4f3bc297ceacb7f016ff951331d352f58e7e791589609ea306/lupa-2.6-cp313-cp313-win32.whl", hash = "sha256:ee9523941ae0a87b5b703417720c5d78f72d2f5bc23883a2ea80a949a3ed9e75", size = 1419558, upload-time = "2025-10-24T07:18:48.371Z" },
+ { url = "https://files.pythonhosted.org/packages/f9/b4/55e885834c847ea610e111d87b9ed4768f0afdaeebc00cd46810f25029f6/lupa-2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b1335a5835b0a25ebdbc75cf0bda195e54d133e4d994877ef025e218c2e59db9", size = 1683424, upload-time = "2025-10-24T07:18:50.976Z" },
+ { url = "https://files.pythonhosted.org/packages/66/9d/d9427394e54d22a35d1139ef12e845fd700d4872a67a34db32516170b746/lupa-2.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:dcb6d0a3264873e1653bc188499f48c1fb4b41a779e315eba45256cfe7bc33c1", size = 953818, upload-time = "2025-10-24T07:18:53.378Z" },
+ { url = "https://files.pythonhosted.org/packages/10/41/27bbe81953fb2f9ecfced5d9c99f85b37964cfaf6aa8453bb11283983721/lupa-2.6-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:a37e01f2128f8c36106726cb9d360bac087d58c54b4522b033cc5691c584db18", size = 1915850, upload-time = "2025-10-24T07:18:55.259Z" },
+ { url = "https://files.pythonhosted.org/packages/a3/98/f9ff60db84a75ba8725506bbf448fb085bc77868a021998ed2a66d920568/lupa-2.6-cp314-cp314-macosx_11_0_x86_64.whl", hash = "sha256:458bd7e9ff3c150b245b0fcfbb9bd2593d1152ea7f0a7b91c1d185846da033fe", size = 982344, upload-time = "2025-10-24T07:18:57.05Z" },
+ { url = "https://files.pythonhosted.org/packages/41/f7/f39e0f1c055c3b887d86b404aaf0ca197b5edfd235a8b81b45b25bac7fc3/lupa-2.6-cp314-cp314-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:052ee82cac5206a02df77119c325339acbc09f5ce66967f66a2e12a0f3211cad", size = 1156543, upload-time = "2025-10-24T07:18:59.251Z" },
+ { url = "https://files.pythonhosted.org/packages/9e/9c/59e6cffa0d672d662ae17bd7ac8ecd2c89c9449dee499e3eb13ca9cd10d9/lupa-2.6-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96594eca3c87dd07938009e95e591e43d554c1dbd0385be03c100367141db5a8", size = 1047974, upload-time = "2025-10-24T07:19:01.449Z" },
+ { url = "https://files.pythonhosted.org/packages/23/c6/a04e9cef7c052717fcb28fb63b3824802488f688391895b618e39be0f684/lupa-2.6-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8faddd9d198688c8884091173a088a8e920ecc96cda2ffed576a23574c4b3f6", size = 2073458, upload-time = "2025-10-24T07:19:03.369Z" },
+ { url = "https://files.pythonhosted.org/packages/e6/10/824173d10f38b51fc77785228f01411b6ca28826ce27404c7c912e0e442c/lupa-2.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:daebb3a6b58095c917e76ba727ab37b27477fb926957c825205fbda431552134", size = 1067683, upload-time = "2025-10-24T07:19:06.2Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/dc/9692fbcf3c924d9c4ece2d8d2f724451ac2e09af0bd2a782db1cef34e799/lupa-2.6-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f3154e68972befe0f81564e37d8142b5d5d79931a18309226a04ec92487d4ea3", size = 1171892, upload-time = "2025-10-24T07:19:08.544Z" },
+ { url = "https://files.pythonhosted.org/packages/84/ff/e318b628d4643c278c96ab3ddea07fc36b075a57383c837f5b11e537ba9d/lupa-2.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e4dadf77b9fedc0bfa53417cc28dc2278a26d4cbd95c29f8927ad4d8fe0a7ef9", size = 2166641, upload-time = "2025-10-24T07:19:10.485Z" },
+ { url = "https://files.pythonhosted.org/packages/12/f7/a6f9ec2806cf2d50826980cdb4b3cffc7691dc6f95e13cc728846d5cb793/lupa-2.6-cp314-cp314-win32.whl", hash = "sha256:cb34169c6fa3bab3e8ac58ca21b8a7102f6a94b6a5d08d3636312f3f02fafd8f", size = 1456857, upload-time = "2025-10-24T07:19:37.989Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/de/df71896f25bdc18360fdfa3b802cd7d57d7fede41a0e9724a4625b412c85/lupa-2.6-cp314-cp314-win_amd64.whl", hash = "sha256:b74f944fe46c421e25d0f8692aef1e842192f6f7f68034201382ac440ef9ea67", size = 1731191, upload-time = "2025-10-24T07:19:40.281Z" },
+ { url = "https://files.pythonhosted.org/packages/47/3c/a1f23b01c54669465f5f4c4083107d496fbe6fb45998771420e9aadcf145/lupa-2.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0e21b716408a21ab65723f8841cf7f2f37a844b7a965eeabb785e27fca4099cf", size = 999343, upload-time = "2025-10-24T07:19:12.519Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/6d/501994291cb640bfa2ccf7f554be4e6914afa21c4026bd01bff9ca8aac57/lupa-2.6-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:589db872a141bfff828340079bbdf3e9a31f2689f4ca0d88f97d9e8c2eae6142", size = 2000730, upload-time = "2025-10-24T07:19:14.869Z" },
+ { url = "https://files.pythonhosted.org/packages/53/a5/457ffb4f3f20469956c2d4c4842a7675e884efc895b2f23d126d23e126cc/lupa-2.6-cp314-cp314t-macosx_11_0_x86_64.whl", hash = "sha256:cd852a91a4a9d4dcbb9a58100f820a75a425703ec3e3f049055f60b8533b7953", size = 1021553, upload-time = "2025-10-24T07:19:17.123Z" },
+ { url = "https://files.pythonhosted.org/packages/51/6b/36bb5a5d0960f2a5c7c700e0819abb76fd9bf9c1d8a66e5106416d6e9b14/lupa-2.6-cp314-cp314t-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:0334753be028358922415ca97a64a3048e4ed155413fc4eaf87dd0a7e2752983", size = 1133275, upload-time = "2025-10-24T07:19:20.51Z" },
+ { url = "https://files.pythonhosted.org/packages/19/86/202ff4429f663013f37d2229f6176ca9f83678a50257d70f61a0a97281bf/lupa-2.6-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:661d895cd38c87658a34780fac54a690ec036ead743e41b74c3fb81a9e65a6aa", size = 1038441, upload-time = "2025-10-24T07:19:22.509Z" },
+ { url = "https://files.pythonhosted.org/packages/a7/42/d8125f8e420714e5b52e9c08d88b5329dfb02dcca731b4f21faaee6cc5b5/lupa-2.6-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aa58454ccc13878cc177c62529a2056be734da16369e451987ff92784994ca7", size = 2058324, upload-time = "2025-10-24T07:19:24.979Z" },
+ { url = "https://files.pythonhosted.org/packages/2b/2c/47bf8b84059876e877a339717ddb595a4a7b0e8740bacae78ba527562e1c/lupa-2.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1425017264e470c98022bba8cff5bd46d054a827f5df6b80274f9cc71dafd24f", size = 1060250, upload-time = "2025-10-24T07:19:27.262Z" },
+ { url = "https://files.pythonhosted.org/packages/c2/06/d88add2b6406ca1bdec99d11a429222837ca6d03bea42ca75afa169a78cb/lupa-2.6-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:224af0532d216e3105f0a127410f12320f7c5f1aa0300bdf9646b8d9afb0048c", size = 1151126, upload-time = "2025-10-24T07:19:29.522Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/a0/89e6a024c3b4485b89ef86881c9d55e097e7cb0bdb74efb746f2fa6a9a76/lupa-2.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9abb98d5a8fd27c8285302e82199f0e56e463066f88f619d6594a450bf269d80", size = 2153693, upload-time = "2025-10-24T07:19:31.379Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/36/a0f007dc58fc1bbf51fb85dcc82fcb1f21b8c4261361de7dab0e3d8521ef/lupa-2.6-cp314-cp314t-win32.whl", hash = "sha256:1849efeba7a8f6fb8aa2c13790bee988fd242ae404bd459509640eeea3d1e291", size = 1590104, upload-time = "2025-10-24T07:19:33.514Z" },
+ { url = "https://files.pythonhosted.org/packages/7d/5e/db903ce9cf82c48d6b91bf6d63ae4c8d0d17958939a4e04ba6b9f38b8643/lupa-2.6-cp314-cp314t-win_amd64.whl", hash = "sha256:fc1498d1a4fc028bc521c26d0fad4ca00ed63b952e32fb95949bda76a04bad52", size = 1913818, upload-time = "2025-10-24T07:19:36.039Z" },
+]
+
[[package]]
name = "markdown-it-py"
version = "4.0.0"
@@ -1052,6 +1486,31 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" },
]
+[[package]]
+name = "mcp"
+version = "1.26.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "anyio" },
+ { name = "httpx" },
+ { name = "httpx-sse" },
+ { name = "jsonschema" },
+ { name = "pydantic" },
+ { name = "pydantic-settings" },
+ { name = "pyjwt", extra = ["crypto"] },
+ { name = "python-multipart" },
+ { name = "pywin32", marker = "sys_platform == 'win32'" },
+ { name = "sse-starlette" },
+ { name = "starlette" },
+ { name = "typing-extensions" },
+ { name = "typing-inspection" },
+ { name = "uvicorn", marker = "sys_platform != 'emscripten'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" },
+]
+
[[package]]
name = "mdurl"
version = "0.1.2"
@@ -1061,6 +1520,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" },
]
+[[package]]
+name = "more-itertools"
+version = "10.8.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ea/5d/38b681d3fce7a266dd9ab73c66959406d565b3e85f21d5e66e1181d93721/more_itertools-10.8.0.tar.gz", hash = "sha256:f638ddf8a1a0d134181275fb5d58b086ead7c6a72429ad725c67503f13ba30bd", size = 137431, upload-time = "2025-09-02T15:23:11.018Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl", hash = "sha256:52d4362373dcf7c52546bc4af9a86ee7c4579df9a8dc268be0a2f949d376cc9b", size = 69667, upload-time = "2025-09-02T15:23:09.635Z" },
+]
+
[[package]]
name = "multidict"
version = "6.7.0"
@@ -1199,15 +1667,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/da/7d22601b625e241d4f23ef1ebff8acfc60da633c9e7e7922e24d10f592b3/multidict-6.7.0-py3-none-any.whl", hash = "sha256:394fc5c42a333c9ffc3e421a4c85e08580d990e08b99f6bf35b4132114c5dcb3", size = 12317, upload-time = "2025-10-06T14:52:29.272Z" },
]
-[[package]]
-name = "mypy-extensions"
-version = "1.1.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" }
-wheels = [
- { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" },
-]
-
[[package]]
name = "openai"
version = "2.13.0"
@@ -1227,6 +1686,31 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/bb/d5/eb52edff49d3d5ea116e225538c118699ddeb7c29fa17ec28af14bc10033/openai-2.13.0-py3-none-any.whl", hash = "sha256:746521065fed68df2f9c2d85613bb50844343ea81f60009b60e6a600c9352c79", size = 1066837, upload-time = "2025-12-16T18:19:43.124Z" },
]
+[[package]]
+name = "openapi-pydantic"
+version = "0.5.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "pydantic" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/02/2e/58d83848dd1a79cb92ed8e63f6ba901ca282c5f09d04af9423ec26c56fd7/openapi_pydantic-0.5.1.tar.gz", hash = "sha256:ff6835af6bde7a459fb93eb93bb92b8749b754fc6e51b2f1590a19dc3005ee0d", size = 60892, upload-time = "2025-01-08T19:29:27.083Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/12/cf/03675d8bd8ecbf4445504d8071adab19f5f993676795708e36402ab38263/openapi_pydantic-0.5.1-py3-none-any.whl", hash = "sha256:a3a09ef4586f5bd760a8df7f43028b60cafb6d9f61de2acba9574766255ab146", size = 96381, upload-time = "2025-01-08T19:29:25.275Z" },
+]
+
+[[package]]
+name = "opentelemetry-api"
+version = "1.39.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "importlib-metadata" },
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/97/b9/3161be15bb8e3ad01be8be5a968a9237c3027c5be504362ff800fca3e442/opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c", size = 65767, upload-time = "2025-12-11T13:32:39.182Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/cf/df/d3f1ddf4bb4cb50ed9b1139cc7b1c54c34a1e7ce8fd1b9a37c0d1551a6bd/opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950", size = 66356, upload-time = "2025-12-11T13:32:17.304Z" },
+]
+
[[package]]
name = "osmosis-ai"
source = { editable = "." }
@@ -1243,18 +1727,22 @@ dependencies = [
[package.optional-dependencies]
dev = [
{ name = "anyio" },
- { name = "black" },
- { name = "isort" },
+ { name = "fastmcp" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
+ { name = "ruff" },
]
full = [
{ name = "fastapi" },
+ { name = "fastmcp" },
{ name = "pyarrow" },
{ name = "pydantic-settings" },
{ name = "rich" },
{ name = "uvicorn" },
]
+mcp = [
+ { name = "fastmcp" },
+]
server = [
{ name = "fastapi" },
{ name = "pyarrow" },
@@ -1266,25 +1754,31 @@ server = [
[package.metadata]
requires-dist = [
{ name = "anyio", marker = "extra == 'dev'", specifier = ">=4.0.0" },
- { name = "black", marker = "extra == 'dev'", specifier = ">=25.0.0,<26.0.0" },
+ { name = "fastapi", marker = "extra == 'full'", specifier = ">=0.100.0,<1.0.0" },
{ name = "fastapi", marker = "extra == 'server'", specifier = ">=0.100.0,<1.0.0" },
+ { name = "fastmcp", marker = "extra == 'dev'", specifier = ">=2.0.0" },
+ { name = "fastmcp", marker = "extra == 'full'", specifier = ">=2.0.0" },
+ { name = "fastmcp", marker = "extra == 'mcp'", specifier = ">=2.0.0" },
{ name = "httpx", specifier = ">=0.25.0,<1.0.0" },
- { name = "isort", marker = "extra == 'dev'", specifier = ">=5.0.0,<7.0.0" },
{ name = "litellm", specifier = ">=1.40.0,<2.0.0" },
- { name = "osmosis-ai", extras = ["server"], marker = "extra == 'full'" },
+ { name = "pyarrow", marker = "extra == 'full'", specifier = ">=14.0.0" },
{ name = "pyarrow", marker = "extra == 'server'", specifier = ">=14.0.0" },
{ name = "pydantic", specifier = ">=2.0.0,<3.0.0" },
+ { name = "pydantic-settings", marker = "extra == 'full'", specifier = ">=2.0.0,<3.0.0" },
{ name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.0.0,<3.0.0" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0,<10.0.0" },
{ name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0,<2.0.0" },
{ name = "python-dotenv", specifier = ">=0.1.0,<2.0.0" },
{ name = "pyyaml", specifier = ">=6.0,<7.0" },
{ name = "requests", specifier = ">=2.0.0,<3.0.0" },
+ { name = "rich", marker = "extra == 'full'", specifier = ">=13.0.0" },
{ name = "rich", marker = "extra == 'server'", specifier = ">=13.0.0" },
+ { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0,<1.0.0" },
{ name = "tqdm", specifier = ">=4.0.0,<5.0.0" },
+ { name = "uvicorn", marker = "extra == 'full'", specifier = ">=0.23.0,<1.0.0" },
{ name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.23.0,<1.0.0" },
]
-provides-extras = ["server", "dev", "full"]
+provides-extras = ["server", "mcp", "dev", "full"]
[[package]]
name = "packaging"
@@ -1296,12 +1790,21 @@ wheels = [
]
[[package]]
-name = "pathspec"
-version = "0.12.1"
+name = "pathable"
+version = "0.4.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/67/93/8f2c2075b180c12c1e9f6a09d1a985bc2036906b13dff1d8917e395f2048/pathable-0.4.4.tar.gz", hash = "sha256:6905a3cd17804edfac7875b5f6c9142a218c7caef78693c2dbbbfbac186d88b2", size = 8124, upload-time = "2025-01-10T18:43:13.247Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7d/eb/b6260b31b1a96386c0a880edebe26f89669098acea8e0318bff6adb378fd/pathable-0.4.4-py3-none-any.whl", hash = "sha256:5ae9e94793b6ef5a4cbe0a7ce9dbbefc1eec38df253763fd0aeeacf2762dbbc2", size = 9592, upload-time = "2025-01-10T18:43:11.88Z" },
+]
+
+[[package]]
+name = "pathvalidate"
+version = "3.3.1"
source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/ca/bc/f35b8446f4531a7cb215605d100cd88b7ac6f44ab3fc94870c120ab3adbf/pathspec-0.12.1.tar.gz", hash = "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712", size = 51043, upload-time = "2023-12-10T22:30:45Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/fa/2a/52a8da6fe965dea6192eb716b357558e103aea0a1e9a8352ad575a8406ca/pathvalidate-3.3.1.tar.gz", hash = "sha256:b18c07212bfead624345bb8e1d6141cdcf15a39736994ea0b94035ad2b1ba177", size = 63262, upload-time = "2025-06-15T09:07:20.736Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", size = 31191, upload-time = "2023-12-10T22:30:43.14Z" },
+ { url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305, upload-time = "2025-06-15T09:07:19.117Z" },
]
[[package]]
@@ -1322,6 +1825,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
]
+[[package]]
+name = "prometheus-client"
+version = "0.24.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f0/58/a794d23feb6b00fc0c72787d7e87d872a6730dd9ed7c7b3e954637d8f280/prometheus_client-0.24.1.tar.gz", hash = "sha256:7e0ced7fbbd40f7b84962d5d2ab6f17ef88a72504dcf7c0b40737b43b2a461f9", size = 85616, upload-time = "2026-01-14T15:26:26.965Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/74/c3/24a2f845e3917201628ecaba4f18bab4d18a337834c1df2a159ee9d22a42/prometheus_client-0.24.1-py3-none-any.whl", hash = "sha256:150db128af71a5c2482b36e588fc8a6b95e498750da4b17065947c16070f4055", size = 64057, upload-time = "2026-01-14T15:26:24.42Z" },
+]
+
[[package]]
name = "propcache"
version = "0.4.1"
@@ -1436,6 +1948,47 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" },
]
+[[package]]
+name = "py-key-value-aio"
+version = "0.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "beartype" },
+ { name = "py-key-value-shared" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/93/ce/3136b771dddf5ac905cc193b461eb67967cf3979688c6696e1f2cdcde7ea/py_key_value_aio-0.3.0.tar.gz", hash = "sha256:858e852fcf6d696d231266da66042d3355a7f9871650415feef9fca7a6cd4155", size = 50801, upload-time = "2025-11-17T16:50:04.711Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/99/10/72f6f213b8f0bce36eff21fda0a13271834e9eeff7f9609b01afdc253c79/py_key_value_aio-0.3.0-py3-none-any.whl", hash = "sha256:1c781915766078bfd608daa769fefb97e65d1d73746a3dfb640460e322071b64", size = 96342, upload-time = "2025-11-17T16:50:03.801Z" },
+]
+
+[package.optional-dependencies]
+disk = [
+ { name = "diskcache" },
+ { name = "pathvalidate" },
+]
+keyring = [
+ { name = "keyring" },
+]
+memory = [
+ { name = "cachetools" },
+]
+redis = [
+ { name = "redis" },
+]
+
+[[package]]
+name = "py-key-value-shared"
+version = "0.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "beartype" },
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7b/e4/1971dfc4620a3a15b4579fe99e024f5edd6e0967a71154771a059daff4db/py_key_value_shared-0.3.0.tar.gz", hash = "sha256:8fdd786cf96c3e900102945f92aa1473138ebe960ef49da1c833790160c28a4b", size = 11666, upload-time = "2025-11-17T16:50:06.849Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/51/e4/b8b0a03ece72f47dce2307d36e1c34725b7223d209fc679315ffe6a4e2c3/py_key_value_shared-0.3.0-py3-none-any.whl", hash = "sha256:5b0efba7ebca08bb158b1e93afc2f07d30b8f40c2fc12ce24a4c0d84f42f9298", size = 19560, upload-time = "2025-11-17T16:50:05.954Z" },
+]
+
[[package]]
name = "pyarrow"
version = "22.0.0"
@@ -1493,6 +2046,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/7b/03/f335d6c52b4a4761bcc83499789a1e2e16d9d201a58c327a9b5cc9a41bd9/pyarrow-22.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0c34fe18094686194f204a3b1787a27456897d8a2d62caf84b61e8dfbc0252ae", size = 29185594, upload-time = "2025-10-24T10:09:53.111Z" },
]
+[[package]]
+name = "pycparser"
+version = "3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" },
+]
+
[[package]]
name = "pydantic"
version = "2.12.5"
@@ -1508,6 +2070,11 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" },
]
+[package.optional-dependencies]
+email = [
+ { name = "email-validator" },
+]
+
[[package]]
name = "pydantic-core"
version = "2.41.5"
@@ -1640,6 +2207,30 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" },
]
+[[package]]
+name = "pydocket"
+version = "0.17.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cloudpickle" },
+ { name = "croniter" },
+ { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
+ { name = "fakeredis", extra = ["lua"] },
+ { name = "opentelemetry-api" },
+ { name = "prometheus-client" },
+ { name = "py-key-value-aio", extra = ["memory", "redis"] },
+ { name = "python-json-logger" },
+ { name = "redis" },
+ { name = "rich" },
+ { name = "taskgroup", marker = "python_full_version < '3.11'" },
+ { name = "typer" },
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/cd/b2/5e12dbe2acf59e4499285e8eee66e8e81b6ba2f553696d2f4ccca0a7978c/pydocket-0.17.7.tar.gz", hash = "sha256:5c77ec6731a167cdcb44174abf793fe63e7b6c1c1c8a799cc6ec7502b361ee77", size = 347071, upload-time = "2026-02-11T21:01:31.744Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/c9/c7/68f2553819965326f968375f02597d49efe71b309ba9d8fef539aeb51c48/pydocket-0.17.7-py3-none-any.whl", hash = "sha256:d1e0921ac02026c4a0140fc72a3848545f3e91e6e74c6e32c588489017c130b2", size = 94608, upload-time = "2026-02-11T21:01:30.111Z" },
+]
+
[[package]]
name = "pygments"
version = "2.19.2"
@@ -1649,6 +2240,29 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" },
]
+[[package]]
+name = "pyjwt"
+version = "2.11.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/5c/5a/b46fa56bf322901eee5b0454a34343cdbdae202cd421775a8ee4e42fd519/pyjwt-2.11.0.tar.gz", hash = "sha256:35f95c1f0fbe5d5ba6e43f00271c275f7a1a4db1dab27bf708073b75318ea623", size = 98019, upload-time = "2026-01-30T19:59:55.694Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/6f/01/c26ce75ba460d5cd503da9e13b21a33804d38c2165dec7b716d06b13010c/pyjwt-2.11.0-py3-none-any.whl", hash = "sha256:94a6bde30eb5c8e04fee991062b534071fd1439ef58d2adc9ccb823e7bcd0469", size = 28224, upload-time = "2026-01-30T19:59:54.539Z" },
+]
+
+[package.optional-dependencies]
+crypto = [
+ { name = "cryptography" },
+]
+
+[[package]]
+name = "pyperclip"
+version = "1.11.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e8/52/d87eba7cb129b81563019d1679026e7a112ef76855d6159d24754dbd2a51/pyperclip-1.11.0.tar.gz", hash = "sha256:244035963e4428530d9e3a6101a1ef97209c6825edab1567beac148ccc1db1b6", size = 12185, upload-time = "2025-09-26T14:40:37.245Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063, upload-time = "2025-09-26T14:40:36.069Z" },
+]
+
[[package]]
name = "pytest"
version = "9.0.2"
@@ -1681,6 +2295,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
]
+[[package]]
+name = "python-dateutil"
+version = "2.9.0.post0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "six" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" },
+]
+
[[package]]
name = "python-dotenv"
version = "1.2.1"
@@ -1691,12 +2317,61 @@ wheels = [
]
[[package]]
-name = "pytokens"
-version = "0.3.0"
+name = "python-json-logger"
+version = "4.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/29/bf/eca6a3d43db1dae7070f70e160ab20b807627ba953663ba07928cdd3dc58/python_json_logger-4.0.0.tar.gz", hash = "sha256:f58e68eb46e1faed27e0f574a55a0455eecd7b8a5b88b85a784519ba3cff047f", size = 17683, upload-time = "2025-10-06T04:15:18.984Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/51/e5/fecf13f06e5e5f67e8837d777d1bc43fac0ed2b77a676804df5c34744727/python_json_logger-4.0.0-py3-none-any.whl", hash = "sha256:af09c9daf6a813aa4cc7180395f50f2a9e5fa056034c9953aec92e381c5ba1e2", size = 15548, upload-time = "2025-10-06T04:15:17.553Z" },
+]
+
+[[package]]
+name = "python-multipart"
+version = "0.0.22"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/01/979e98d542a70714b0cb2b6728ed0b7c46792b695e3eaec3e20711271ca3/python_multipart-0.0.22.tar.gz", hash = "sha256:7340bef99a7e0032613f56dc36027b959fd3b30a787ed62d310e951f7c3a3a58", size = 37612, upload-time = "2026-01-25T10:15:56.219Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/1b/d0/397f9626e711ff749a95d96b7af99b9c566a9bb5129b8e4c10fc4d100304/python_multipart-0.0.22-py3-none-any.whl", hash = "sha256:2b2cd894c83d21bf49d702499531c7bafd057d730c201782048f7945d82de155", size = 24579, upload-time = "2026-01-25T10:15:54.811Z" },
+]
+
+[[package]]
+name = "pytz"
+version = "2025.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f8/bf/abbd3cdfb8fbc7fb3d4d38d320f2441b1e7cbe29be4f23797b4a2b5d8aac/pytz-2025.2.tar.gz", hash = "sha256:360b9e3dbb49a209c21ad61809c7fb453643e048b38924c765813546746e81c3", size = 320884, upload-time = "2025-03-25T02:25:00.538Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/81/c4/34e93fe5f5429d7570ec1fa436f1986fb1f00c3e0f43a589fe2bbcd22c3f/pytz-2025.2-py2.py3-none-any.whl", hash = "sha256:5ddf76296dd8c44c26eb8f4b6f35488f3ccbf6fbbd7adee0b7262d43f0ec2f00", size = 509225, upload-time = "2025-03-25T02:24:58.468Z" },
+]
+
+[[package]]
+name = "pywin32"
+version = "311"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7b/40/44efbb0dfbd33aca6a6483191dae0716070ed99e2ecb0c53683f400a0b4f/pywin32-311-cp310-cp310-win32.whl", hash = "sha256:d03ff496d2a0cd4a5893504789d4a15399133fe82517455e78bad62efbb7f0a3", size = 8760432, upload-time = "2025-07-14T20:13:05.9Z" },
+ { url = "https://files.pythonhosted.org/packages/5e/bf/360243b1e953bd254a82f12653974be395ba880e7ec23e3731d9f73921cc/pywin32-311-cp310-cp310-win_amd64.whl", hash = "sha256:797c2772017851984b97180b0bebe4b620bb86328e8a884bb626156295a63b3b", size = 9590103, upload-time = "2025-07-14T20:13:07.698Z" },
+ { url = "https://files.pythonhosted.org/packages/57/38/d290720e6f138086fb3d5ffe0b6caa019a791dd57866940c82e4eeaf2012/pywin32-311-cp310-cp310-win_arm64.whl", hash = "sha256:0502d1facf1fed4839a9a51ccbcc63d952cf318f78ffc00a7e78528ac27d7a2b", size = 8778557, upload-time = "2025-07-14T20:13:11.11Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/af/449a6a91e5d6db51420875c54f6aff7c97a86a3b13a0b4f1a5c13b988de3/pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151", size = 8697031, upload-time = "2025-07-14T20:13:13.266Z" },
+ { url = "https://files.pythonhosted.org/packages/51/8f/9bb81dd5bb77d22243d33c8397f09377056d5c687aa6d4042bea7fbf8364/pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503", size = 9508308, upload-time = "2025-07-14T20:13:15.147Z" },
+ { url = "https://files.pythonhosted.org/packages/44/7b/9c2ab54f74a138c491aba1b1cd0795ba61f144c711daea84a88b63dc0f6c/pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2", size = 8703930, upload-time = "2025-07-14T20:13:16.945Z" },
+ { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" },
+ { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" },
+ { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/be/3fd5de0979fcb3994bfee0d65ed8ca9506a8a1260651b86174f6a86f52b3/pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d", size = 8705700, upload-time = "2025-07-14T20:13:26.471Z" },
+ { url = "https://files.pythonhosted.org/packages/e3/28/e0a1909523c6890208295a29e05c2adb2126364e289826c0a8bc7297bd5c/pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d", size = 9494700, upload-time = "2025-07-14T20:13:28.243Z" },
+ { url = "https://files.pythonhosted.org/packages/04/bf/90339ac0f55726dce7d794e6d79a18a91265bdf3aa70b6b9ca52f35e022a/pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a", size = 8709318, upload-time = "2025-07-14T20:13:30.348Z" },
+ { url = "https://files.pythonhosted.org/packages/c9/31/097f2e132c4f16d99a22bfb777e0fd88bd8e1c634304e102f313af69ace5/pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee", size = 8840714, upload-time = "2025-07-14T20:13:32.449Z" },
+ { url = "https://files.pythonhosted.org/packages/90/4b/07c77d8ba0e01349358082713400435347df8426208171ce297da32c313d/pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87", size = 9656800, upload-time = "2025-07-14T20:13:34.312Z" },
+ { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" },
+]
+
+[[package]]
+name = "pywin32-ctypes"
+version = "0.2.3"
source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/4e/8d/a762be14dae1c3bf280202ba3172020b2b0b4c537f94427435f19c413b72/pytokens-0.3.0.tar.gz", hash = "sha256:2f932b14ed08de5fcf0b391ace2642f858f1394c0857202959000b68ed7a458a", size = 17644, upload-time = "2025-11-05T13:36:35.34Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/85/9f/01a1a99704853cb63f253eea009390c88e7131c67e66a0a02099a8c917cb/pywin32-ctypes-0.2.3.tar.gz", hash = "sha256:d162dc04946d704503b2edc4d55f3dba5c1d539ead017afa00142c38b9885755", size = 29471, upload-time = "2024-08-14T10:15:34.626Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/84/25/d9db8be44e205a124f6c98bc0324b2bb149b7431c53877fc6d1038dddaf5/pytokens-0.3.0-py3-none-any.whl", hash = "sha256:95b2b5eaf832e469d141a378872480ede3f251a5a5041b8ec6e581d3ac71bbf3", size = 12195, upload-time = "2025-11-05T13:36:33.183Z" },
+ { url = "https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl", hash = "sha256:8a1513379d709975552d202d942d9837758905c8d01eb82b8bcc30918929e7b8", size = 30756, upload-time = "2024-08-14T10:15:33.187Z" },
]
[[package]]
@@ -1763,18 +2438,30 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
]
+[[package]]
+name = "redis"
+version = "7.2.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "async-timeout", marker = "python_full_version < '3.11.3'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9f/32/6fac13a11e73e1bc67a2ae821a72bfe4c2d8c4c48f0267e4a952be0f1bae/redis-7.2.0.tar.gz", hash = "sha256:4dd5bf4bd4ae80510267f14185a15cba2a38666b941aff68cccf0256b51c1f26", size = 4901247, upload-time = "2026-02-16T17:16:22.797Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/86/cf/f6180b67f99688d83e15c84c5beda831d1d341e95872d224f87ccafafe61/redis-7.2.0-py3-none-any.whl", hash = "sha256:01f591f8598e483f1842d429e8ae3a820804566f1c73dca1b80e23af9fba0497", size = 394898, upload-time = "2026-02-16T17:16:20.693Z" },
+]
+
[[package]]
name = "referencing"
-version = "0.37.0"
+version = "0.36.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "attrs" },
{ name = "rpds-py" },
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/2f/db/98b5c277be99dd18bfd91dd04e1b759cad18d1a338188c936e92f921c7e2/referencing-0.36.2.tar.gz", hash = "sha256:df2e89862cd09deabbdba16944cc3f10feb6b3e6f18e902f7cc25609a34775aa", size = 74744, upload-time = "2025-01-25T08:48:16.138Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl", hash = "sha256:381329a9f99628c9069361716891d34ad94af76e461dcb0335825aecc7692231", size = 26766, upload-time = "2025-10-13T15:30:47.625Z" },
+ { url = "https://files.pythonhosted.org/packages/c1/b1/3baf80dc6d2b7bc27a95a67752d0208e410351e3feb4eb78de5f77454d8d/referencing-0.36.2-py3-none-any.whl", hash = "sha256:e8699adbbf8b5c7de96d8ffa0eb5c158b3beafce084968e2ea8bb08c6794dcd0", size = 26775, upload-time = "2025-01-25T08:48:14.241Z" },
]
[[package]]
@@ -1912,6 +2599,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/25/7a/b0178788f8dc6cafce37a212c99565fa1fe7872c70c6c9c1e1a372d9d88f/rich-14.2.0-py3-none-any.whl", hash = "sha256:76bc51fe2e57d2b1be1f96c524b890b816e334ab4c1e45888799bfaab0021edd", size = 243393, upload-time = "2025-10-09T14:16:51.245Z" },
]
+[[package]]
+name = "rich-rst"
+version = "1.3.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "docutils" },
+ { name = "rich" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/bc/6d/a506aaa4a9eaa945ed8ab2b7347859f53593864289853c5d6d62b77246e0/rich_rst-1.3.2.tar.gz", hash = "sha256:a1196fdddf1e364b02ec68a05e8ff8f6914fee10fbca2e6b6735f166bb0da8d4", size = 14936, upload-time = "2025-10-14T16:49:45.332Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/13/2f/b4530fbf948867702d0a3f27de4a6aab1d156f406d72852ab902c4d04de9/rich_rst-1.3.2-py3-none-any.whl", hash = "sha256:a99b4907cbe118cf9d18b0b44de272efa61f15117c61e39ebdc431baf5df722a", size = 12567, upload-time = "2025-10-14T16:49:42.953Z" },
+]
+
[[package]]
name = "rpds-py"
version = "0.30.0"
@@ -2034,6 +2734,44 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/b7/b95708304cd49b7b6f82fdd039f1748b66ec2b21d6a45180910802f1abf1/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:ac37f9f516c51e5753f27dfdef11a88330f04de2d564be3991384b2f3535d02e", size = 562191, upload-time = "2025-11-30T20:24:36.853Z" },
]
+[[package]]
+name = "ruff"
+version = "0.15.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/04/dc/4e6ac71b511b141cf626357a3946679abeba4cf67bc7cc5a17920f31e10d/ruff-0.15.1.tar.gz", hash = "sha256:c590fe13fb57c97141ae975c03a1aedb3d3156030cabd740d6ff0b0d601e203f", size = 4540855, upload-time = "2026-02-12T23:09:09.998Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/23/bf/e6e4324238c17f9d9120a9d60aa99a7daaa21204c07fcd84e2ef03bb5fd1/ruff-0.15.1-py3-none-linux_armv6l.whl", hash = "sha256:b101ed7cf4615bda6ffe65bdb59f964e9f4a0d3f85cbf0e54f0ab76d7b90228a", size = 10367819, upload-time = "2026-02-12T23:09:03.598Z" },
+ { url = "https://files.pythonhosted.org/packages/b3/ea/c8f89d32e7912269d38c58f3649e453ac32c528f93bb7f4219258be2e7ed/ruff-0.15.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:939c995e9277e63ea632cc8d3fae17aa758526f49a9a850d2e7e758bfef46602", size = 10798618, upload-time = "2026-02-12T23:09:22.928Z" },
+ { url = "https://files.pythonhosted.org/packages/5e/0f/1d0d88bc862624247d82c20c10d4c0f6bb2f346559d8af281674cf327f15/ruff-0.15.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1d83466455fdefe60b8d9c8df81d3c1bbb2115cede53549d3b522ce2bc703899", size = 10148518, upload-time = "2026-02-12T23:08:58.339Z" },
+ { url = "https://files.pythonhosted.org/packages/f5/c8/291c49cefaa4a9248e986256df2ade7add79388fe179e0691be06fae6f37/ruff-0.15.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a9457e3c3291024866222b96108ab2d8265b477e5b1534c7ddb1810904858d16", size = 10518811, upload-time = "2026-02-12T23:09:31.865Z" },
+ { url = "https://files.pythonhosted.org/packages/c3/1a/f5707440e5ae43ffa5365cac8bbb91e9665f4a883f560893829cf16a606b/ruff-0.15.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:92c92b003e9d4f7fbd33b1867bb15a1b785b1735069108dfc23821ba045b29bc", size = 10196169, upload-time = "2026-02-12T23:09:17.306Z" },
+ { url = "https://files.pythonhosted.org/packages/2a/ff/26ddc8c4da04c8fd3ee65a89c9fb99eaa5c30394269d424461467be2271f/ruff-0.15.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fe5c41ab43e3a06778844c586251eb5a510f67125427625f9eb2b9526535779", size = 10990491, upload-time = "2026-02-12T23:09:25.503Z" },
+ { url = "https://files.pythonhosted.org/packages/fc/00/50920cb385b89413f7cdb4bb9bc8fc59c1b0f30028d8bccc294189a54955/ruff-0.15.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:66a6dd6df4d80dc382c6484f8ce1bcceb55c32e9f27a8b94c32f6c7331bf14fb", size = 11843280, upload-time = "2026-02-12T23:09:19.88Z" },
+ { url = "https://files.pythonhosted.org/packages/5d/6d/2f5cad8380caf5632a15460c323ae326f1e1a2b5b90a6ee7519017a017ca/ruff-0.15.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6a4a42cbb8af0bda9bcd7606b064d7c0bc311a88d141d02f78920be6acb5aa83", size = 11274336, upload-time = "2026-02-12T23:09:14.907Z" },
+ { url = "https://files.pythonhosted.org/packages/a3/1d/5f56cae1d6c40b8a318513599b35ea4b075d7dc1cd1d04449578c29d1d75/ruff-0.15.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4ab064052c31dddada35079901592dfba2e05f5b1e43af3954aafcbc1096a5b2", size = 11137288, upload-time = "2026-02-12T23:09:07.475Z" },
+ { url = "https://files.pythonhosted.org/packages/cd/20/6f8d7d8f768c93b0382b33b9306b3b999918816da46537d5a61635514635/ruff-0.15.1-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:5631c940fe9fe91f817a4c2ea4e81f47bee3ca4aa646134a24374f3c19ad9454", size = 11070681, upload-time = "2026-02-12T23:08:55.43Z" },
+ { url = "https://files.pythonhosted.org/packages/9a/67/d640ac76069f64cdea59dba02af2e00b1fa30e2103c7f8d049c0cff4cafd/ruff-0.15.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:68138a4ba184b4691ccdc39f7795c66b3c68160c586519e7e8444cf5a53e1b4c", size = 10486401, upload-time = "2026-02-12T23:09:27.927Z" },
+ { url = "https://files.pythonhosted.org/packages/65/3d/e1429f64a3ff89297497916b88c32a5cc88eeca7e9c787072d0e7f1d3e1e/ruff-0.15.1-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:518f9af03bfc33c03bdb4cb63fabc935341bb7f54af500f92ac309ecfbba6330", size = 10197452, upload-time = "2026-02-12T23:09:12.147Z" },
+ { url = "https://files.pythonhosted.org/packages/78/83/e2c3bade17dad63bf1e1c2ffaf11490603b760be149e1419b07049b36ef2/ruff-0.15.1-py3-none-musllinux_1_2_i686.whl", hash = "sha256:da79f4d6a826caaea95de0237a67e33b81e6ec2e25fc7e1993a4015dffca7c61", size = 10693900, upload-time = "2026-02-12T23:09:34.418Z" },
+ { url = "https://files.pythonhosted.org/packages/a1/27/fdc0e11a813e6338e0706e8b39bb7a1d61ea5b36873b351acee7e524a72a/ruff-0.15.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3dd86dccb83cd7d4dcfac303ffc277e6048600dfc22e38158afa208e8bf94a1f", size = 11227302, upload-time = "2026-02-12T23:09:36.536Z" },
+ { url = "https://files.pythonhosted.org/packages/f6/58/ac864a75067dcbd3b95be5ab4eb2b601d7fbc3d3d736a27e391a4f92a5c1/ruff-0.15.1-py3-none-win32.whl", hash = "sha256:660975d9cb49b5d5278b12b03bb9951d554543a90b74ed5d366b20e2c57c2098", size = 10462555, upload-time = "2026-02-12T23:09:29.899Z" },
+ { url = "https://files.pythonhosted.org/packages/e0/5e/d4ccc8a27ecdb78116feac4935dfc39d1304536f4296168f91ed3ec00cd2/ruff-0.15.1-py3-none-win_amd64.whl", hash = "sha256:c820fef9dd5d4172a6570e5721704a96c6679b80cf7be41659ed439653f62336", size = 11599956, upload-time = "2026-02-12T23:09:01.157Z" },
+ { url = "https://files.pythonhosted.org/packages/2a/07/5bda6a85b220c64c65686bc85bd0bbb23b29c62b3a9f9433fa55f17cda93/ruff-0.15.1-py3-none-win_arm64.whl", hash = "sha256:5ff7d5f0f88567850f45081fac8f4ec212be8d0b963e385c3f7d0d2eb4899416", size = 10874604, upload-time = "2026-02-12T23:09:05.515Z" },
+]
+
+[[package]]
+name = "secretstorage"
+version = "3.5.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cryptography" },
+ { name = "jeepney" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" },
+]
+
[[package]]
name = "shellingham"
version = "1.5.4"
@@ -2043,6 +2781,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" },
]
+[[package]]
+name = "six"
+version = "1.17.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" },
+]
+
[[package]]
name = "sniffio"
version = "1.3.1"
@@ -2052,6 +2799,28 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" },
]
+[[package]]
+name = "sortedcontainers"
+version = "2.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e8/c4/ba2f8066cceb6f23394729afe52f3bf7adec04bf9ed2c820b39e19299111/sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88", size = 30594, upload-time = "2021-05-16T22:03:42.897Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/32/46/9cb0e58b2deb7f82b84065f37f3bffeb12413f947f9388e4cac22c4621ce/sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0", size = 29575, upload-time = "2021-05-16T22:03:41.177Z" },
+]
+
+[[package]]
+name = "sse-starlette"
+version = "3.2.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "anyio" },
+ { name = "starlette" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/8b/8d/00d280c03ffd39aaee0e86ec81e2d3b9253036a0f93f51d10503adef0e65/sse_starlette-3.2.0.tar.gz", hash = "sha256:8127594edfb51abe44eac9c49e59b0b01f1039d0c7461c6fd91d4e03b70da422", size = 27253, upload-time = "2026-01-17T13:11:05.62Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/96/7f/832f015020844a8b8f7a9cbc103dd76ba8e3875004c41e08440ea3a2b41a/sse_starlette-3.2.0-py3-none-any.whl", hash = "sha256:5876954bd51920fc2cd51baee47a080eb88a37b5b784e615abb0b283f801cdbf", size = 12763, upload-time = "2026-01-17T13:11:03.775Z" },
+]
+
[[package]]
name = "starlette"
version = "0.50.0"
@@ -2065,6 +2834,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d9/52/1064f510b141bd54025f9b55105e26d1fa970b9be67ad766380a3c9b74b0/starlette-0.50.0-py3-none-any.whl", hash = "sha256:9e5391843ec9b6e472eed1365a78c8098cfceb7a74bfd4d6b1c0c0095efb3bca", size = 74033, upload-time = "2025-11-01T15:25:25.461Z" },
]
+[[package]]
+name = "taskgroup"
+version = "0.2.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "exceptiongroup", marker = "python_full_version < '3.14'" },
+ { name = "typing-extensions", marker = "python_full_version < '3.14'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f0/8d/e218e0160cc1b692e6e0e5ba34e8865dbb171efeb5fc9a704544b3020605/taskgroup-0.2.2.tar.gz", hash = "sha256:078483ac3e78f2e3f973e2edbf6941374fbea81b9c5d0a96f51d297717f4752d", size = 11504, upload-time = "2025-01-03T09:24:13.761Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d1/b1/74babcc824a57904e919f3af16d86c08b524c0691504baf038ef2d7f655c/taskgroup-0.2.2-py2.py3-none-any.whl", hash = "sha256:e2c53121609f4ae97303e9ea1524304b4de6faf9eb2c9280c7f87976479a52fb", size = 14237, upload-time = "2025-01-03T09:24:11.41Z" },
+]
+
[[package]]
name = "tiktoken"
version = "0.12.0"
@@ -2217,6 +2999,21 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" },
]
+[[package]]
+name = "typer"
+version = "0.23.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "annotated-doc" },
+ { name = "click" },
+ { name = "rich" },
+ { name = "shellingham" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/fd/07/b822e1b307d40e263e8253d2384cf98c51aa2368cc7ba9a07e523a1d964b/typer-0.23.1.tar.gz", hash = "sha256:2070374e4d31c83e7b61362fd859aa683576432fd5b026b060ad6b4cd3b86134", size = 120047, upload-time = "2026-02-13T10:04:30.984Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d5/91/9b286ab899c008c2cb05e8be99814807e7fbbd33f0c0c960470826e5ac82/typer-0.23.1-py3-none-any.whl", hash = "sha256:3291ad0d3c701cbf522012faccfbb29352ff16ad262db2139e6b01f15781f14e", size = 56813, upload-time = "2026-02-13T10:04:32.008Z" },
+]
+
[[package]]
name = "typer-slim"
version = "0.21.1"
@@ -2274,6 +3071,74 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ee/d9/d88e73ca598f4f6ff671fb5fde8a32925c2e08a637303a1d12883c7305fa/uvicorn-0.38.0-py3-none-any.whl", hash = "sha256:48c0afd214ceb59340075b4a052ea1ee91c16fbc2a9b1469cca0e54566977b02", size = 68109, upload-time = "2025-10-18T13:46:42.958Z" },
]
+[[package]]
+name = "websockets"
+version = "16.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/20/74/221f58decd852f4b59cc3354cccaf87e8ef695fede361d03dc9a7396573b/websockets-16.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:04cdd5d2d1dacbad0a7bf36ccbcd3ccd5a30ee188f2560b7a62a30d14107b31a", size = 177343, upload-time = "2026-01-10T09:22:21.28Z" },
+ { url = "https://files.pythonhosted.org/packages/19/0f/22ef6107ee52ab7f0b710d55d36f5a5d3ef19e8a205541a6d7ffa7994e5a/websockets-16.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8ff32bb86522a9e5e31439a58addbb0166f0204d64066fb955265c4e214160f0", size = 175021, upload-time = "2026-01-10T09:22:22.696Z" },
+ { url = "https://files.pythonhosted.org/packages/10/40/904a4cb30d9b61c0e278899bf36342e9b0208eb3c470324a9ecbaac2a30f/websockets-16.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:583b7c42688636f930688d712885cf1531326ee05effd982028212ccc13e5957", size = 175320, upload-time = "2026-01-10T09:22:23.94Z" },
+ { url = "https://files.pythonhosted.org/packages/9d/2f/4b3ca7e106bc608744b1cdae041e005e446124bebb037b18799c2d356864/websockets-16.0-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:7d837379b647c0c4c2355c2499723f82f1635fd2c26510e1f587d89bc2199e72", size = 183815, upload-time = "2026-01-10T09:22:25.469Z" },
+ { url = "https://files.pythonhosted.org/packages/86/26/d40eaa2a46d4302becec8d15b0fc5e45bdde05191e7628405a19cf491ccd/websockets-16.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df57afc692e517a85e65b72e165356ed1df12386ecb879ad5693be08fac65dde", size = 185054, upload-time = "2026-01-10T09:22:27.101Z" },
+ { url = "https://files.pythonhosted.org/packages/b0/ba/6500a0efc94f7373ee8fefa8c271acdfd4dca8bd49a90d4be7ccabfc397e/websockets-16.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:2b9f1e0d69bc60a4a87349d50c09a037a2607918746f07de04df9e43252c77a3", size = 184565, upload-time = "2026-01-10T09:22:28.293Z" },
+ { url = "https://files.pythonhosted.org/packages/04/b4/96bf2cee7c8d8102389374a2616200574f5f01128d1082f44102140344cc/websockets-16.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:335c23addf3d5e6a8633f9f8eda77efad001671e80b95c491dd0924587ece0b3", size = 183848, upload-time = "2026-01-10T09:22:30.394Z" },
+ { url = "https://files.pythonhosted.org/packages/02/8e/81f40fb00fd125357814e8c3025738fc4ffc3da4b6b4a4472a82ba304b41/websockets-16.0-cp310-cp310-win32.whl", hash = "sha256:37b31c1623c6605e4c00d466c9d633f9b812ea430c11c8a278774a1fde1acfa9", size = 178249, upload-time = "2026-01-10T09:22:32.083Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/5f/7e40efe8df57db9b91c88a43690ac66f7b7aa73a11aa6a66b927e44f26fa/websockets-16.0-cp310-cp310-win_amd64.whl", hash = "sha256:8e1dab317b6e77424356e11e99a432b7cb2f3ec8c5ab4dabbcee6add48f72b35", size = 178685, upload-time = "2026-01-10T09:22:33.345Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/db/de907251b4ff46ae804ad0409809504153b3f30984daf82a1d84a9875830/websockets-16.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:31a52addea25187bde0797a97d6fc3d2f92b6f72a9370792d65a6e84615ac8a8", size = 177340, upload-time = "2026-01-10T09:22:34.539Z" },
+ { url = "https://files.pythonhosted.org/packages/f3/fa/abe89019d8d8815c8781e90d697dec52523fb8ebe308bf11664e8de1877e/websockets-16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:417b28978cdccab24f46400586d128366313e8a96312e4b9362a4af504f3bbad", size = 175022, upload-time = "2026-01-10T09:22:36.332Z" },
+ { url = "https://files.pythonhosted.org/packages/58/5d/88ea17ed1ded2079358b40d31d48abe90a73c9e5819dbcde1606e991e2ad/websockets-16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af80d74d4edfa3cb9ed973a0a5ba2b2a549371f8a741e0800cb07becdd20f23d", size = 175319, upload-time = "2026-01-10T09:22:37.602Z" },
+ { url = "https://files.pythonhosted.org/packages/d2/ae/0ee92b33087a33632f37a635e11e1d99d429d3d323329675a6022312aac2/websockets-16.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:08d7af67b64d29823fed316505a89b86705f2b7981c07848fb5e3ea3020c1abe", size = 184631, upload-time = "2026-01-10T09:22:38.789Z" },
+ { url = "https://files.pythonhosted.org/packages/c8/c5/27178df583b6c5b31b29f526ba2da5e2f864ecc79c99dae630a85d68c304/websockets-16.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7be95cfb0a4dae143eaed2bcba8ac23f4892d8971311f1b06f3c6b78952ee70b", size = 185870, upload-time = "2026-01-10T09:22:39.893Z" },
+ { url = "https://files.pythonhosted.org/packages/87/05/536652aa84ddc1c018dbb7e2c4cbcd0db884580bf8e95aece7593fde526f/websockets-16.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d6297ce39ce5c2e6feb13c1a996a2ded3b6832155fcfc920265c76f24c7cceb5", size = 185361, upload-time = "2026-01-10T09:22:41.016Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/e2/d5332c90da12b1e01f06fb1b85c50cfc489783076547415bf9f0a659ec19/websockets-16.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1c1b30e4f497b0b354057f3467f56244c603a79c0d1dafce1d16c283c25f6e64", size = 184615, upload-time = "2026-01-10T09:22:42.442Z" },
+ { url = "https://files.pythonhosted.org/packages/77/fb/d3f9576691cae9253b51555f841bc6600bf0a983a461c79500ace5a5b364/websockets-16.0-cp311-cp311-win32.whl", hash = "sha256:5f451484aeb5cafee1ccf789b1b66f535409d038c56966d6101740c1614b86c6", size = 178246, upload-time = "2026-01-10T09:22:43.654Z" },
+ { url = "https://files.pythonhosted.org/packages/54/67/eaff76b3dbaf18dcddabc3b8c1dba50b483761cccff67793897945b37408/websockets-16.0-cp311-cp311-win_amd64.whl", hash = "sha256:8d7f0659570eefb578dacde98e24fb60af35350193e4f56e11190787bee77dac", size = 178684, upload-time = "2026-01-10T09:22:44.941Z" },
+ { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" },
+ { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" },
+ { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" },
+ { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" },
+ { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" },
+ { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" },
+ { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" },
+ { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" },
+ { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" },
+ { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" },
+ { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" },
+ { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" },
+ { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" },
+ { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" },
+ { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" },
+ { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" },
+ { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" },
+ { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" },
+ { url = "https://files.pythonhosted.org/packages/f3/1d/e88022630271f5bd349ed82417136281931e558d628dd52c4d8621b4a0b2/websockets-16.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8cc451a50f2aee53042ac52d2d053d08bf89bcb31ae799cb4487587661c038a0", size = 177406, upload-time = "2026-01-10T09:23:12.178Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/78/e63be1bf0724eeb4616efb1ae1c9044f7c3953b7957799abb5915bffd38e/websockets-16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:daa3b6ff70a9241cf6c7fc9e949d41232d9d7d26fd3522b1ad2b4d62487e9904", size = 175085, upload-time = "2026-01-10T09:23:13.511Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/f4/d3c9220d818ee955ae390cf319a7c7a467beceb24f05ee7aaaa2414345ba/websockets-16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd3cb4adb94a2a6e2b7c0d8d05cb94e6f1c81a0cf9dc2694fb65c7e8d94c42e4", size = 175328, upload-time = "2026-01-10T09:23:14.727Z" },
+ { url = "https://files.pythonhosted.org/packages/63/bc/d3e208028de777087e6fb2b122051a6ff7bbcca0d6df9d9c2bf1dd869ae9/websockets-16.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:781caf5e8eee67f663126490c2f96f40906594cb86b408a703630f95550a8c3e", size = 185044, upload-time = "2026-01-10T09:23:15.939Z" },
+ { url = "https://files.pythonhosted.org/packages/ad/6e/9a0927ac24bd33a0a9af834d89e0abc7cfd8e13bed17a86407a66773cc0e/websockets-16.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caab51a72c51973ca21fa8a18bd8165e1a0183f1ac7066a182ff27107b71e1a4", size = 186279, upload-time = "2026-01-10T09:23:17.148Z" },
+ { url = "https://files.pythonhosted.org/packages/b9/ca/bf1c68440d7a868180e11be653c85959502efd3a709323230314fda6e0b3/websockets-16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19c4dc84098e523fd63711e563077d39e90ec6702aff4b5d9e344a60cb3c0cb1", size = 185711, upload-time = "2026-01-10T09:23:18.372Z" },
+ { url = "https://files.pythonhosted.org/packages/c4/f8/fdc34643a989561f217bb477cbc47a3a07212cbda91c0e4389c43c296ebf/websockets-16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a5e18a238a2b2249c9a9235466b90e96ae4795672598a58772dd806edc7ac6d3", size = 184982, upload-time = "2026-01-10T09:23:19.652Z" },
+ { url = "https://files.pythonhosted.org/packages/dd/d1/574fa27e233764dbac9c52730d63fcf2823b16f0856b3329fc6268d6ae4f/websockets-16.0-cp314-cp314-win32.whl", hash = "sha256:a069d734c4a043182729edd3e9f247c3b2a4035415a9172fd0f1b71658a320a8", size = 177915, upload-time = "2026-01-10T09:23:21.458Z" },
+ { url = "https://files.pythonhosted.org/packages/8a/f1/ae6b937bf3126b5134ce1f482365fde31a357c784ac51852978768b5eff4/websockets-16.0-cp314-cp314-win_amd64.whl", hash = "sha256:c0ee0e63f23914732c6d7e0cce24915c48f3f1512ec1d079ed01fc629dab269d", size = 178381, upload-time = "2026-01-10T09:23:22.715Z" },
+ { url = "https://files.pythonhosted.org/packages/06/9b/f791d1db48403e1f0a27577a6beb37afae94254a8c6f08be4a23e4930bc0/websockets-16.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:a35539cacc3febb22b8f4d4a99cc79b104226a756aa7400adc722e83b0d03244", size = 177737, upload-time = "2026-01-10T09:23:24.523Z" },
+ { url = "https://files.pythonhosted.org/packages/bd/40/53ad02341fa33b3ce489023f635367a4ac98b73570102ad2cdd770dacc9a/websockets-16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b784ca5de850f4ce93ec85d3269d24d4c82f22b7212023c974c401d4980ebc5e", size = 175268, upload-time = "2026-01-10T09:23:25.781Z" },
+ { url = "https://files.pythonhosted.org/packages/74/9b/6158d4e459b984f949dcbbb0c5d270154c7618e11c01029b9bbd1bb4c4f9/websockets-16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:569d01a4e7fba956c5ae4fc988f0d4e187900f5497ce46339c996dbf24f17641", size = 175486, upload-time = "2026-01-10T09:23:27.033Z" },
+ { url = "https://files.pythonhosted.org/packages/e5/2d/7583b30208b639c8090206f95073646c2c9ffd66f44df967981a64f849ad/websockets-16.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:50f23cdd8343b984957e4077839841146f67a3d31ab0d00e6b824e74c5b2f6e8", size = 185331, upload-time = "2026-01-10T09:23:28.259Z" },
+ { url = "https://files.pythonhosted.org/packages/45/b0/cce3784eb519b7b5ad680d14b9673a31ab8dcb7aad8b64d81709d2430aa8/websockets-16.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:152284a83a00c59b759697b7f9e9cddf4e3c7861dd0d964b472b70f78f89e80e", size = 186501, upload-time = "2026-01-10T09:23:29.449Z" },
+ { url = "https://files.pythonhosted.org/packages/19/60/b8ebe4c7e89fb5f6cdf080623c9d92789a53636950f7abacfc33fe2b3135/websockets-16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bc59589ab64b0022385f429b94697348a6a234e8ce22544e3681b2e9331b5944", size = 186062, upload-time = "2026-01-10T09:23:31.368Z" },
+ { url = "https://files.pythonhosted.org/packages/88/a8/a080593f89b0138b6cba1b28f8df5673b5506f72879322288b031337c0b8/websockets-16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32da954ffa2814258030e5a57bc73a3635463238e797c7375dc8091327434206", size = 185356, upload-time = "2026-01-10T09:23:32.627Z" },
+ { url = "https://files.pythonhosted.org/packages/c2/b6/b9afed2afadddaf5ebb2afa801abf4b0868f42f8539bfe4b071b5266c9fe/websockets-16.0-cp314-cp314t-win32.whl", hash = "sha256:5a4b4cc550cb665dd8a47f868c8d04c8230f857363ad3c9caf7a0c3bf8c61ca6", size = 178085, upload-time = "2026-01-10T09:23:33.816Z" },
+ { url = "https://files.pythonhosted.org/packages/9f/3e/28135a24e384493fa804216b79a6a6759a38cc4ff59118787b9fb693df93/websockets-16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b14dc141ed6d2dde437cddb216004bcac6a1df0935d79656387bd41632ba0bbd", size = 178531, upload-time = "2026-01-10T09:23:35.016Z" },
+ { url = "https://files.pythonhosted.org/packages/72/07/c98a68571dcf256e74f1f816b8cc5eae6eb2d3d5cfa44d37f801619d9166/websockets-16.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:349f83cd6c9a415428ee1005cadb5c2c56f4389bc06a9af16103c3bc3dcc8b7d", size = 174947, upload-time = "2026-01-10T09:23:36.166Z" },
+ { url = "https://files.pythonhosted.org/packages/7e/52/93e166a81e0305b33fe416338be92ae863563fe7bce446b0f687b9df5aea/websockets-16.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:4a1aba3340a8dca8db6eb5a7986157f52eb9e436b74813764241981ca4888f03", size = 175260, upload-time = "2026-01-10T09:23:37.409Z" },
+ { url = "https://files.pythonhosted.org/packages/56/0c/2dbf513bafd24889d33de2ff0368190a0e69f37bcfa19009ef819fe4d507/websockets-16.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f4a32d1bd841d4bcbffdcb3d2ce50c09c3909fbead375ab28d0181af89fd04da", size = 176071, upload-time = "2026-01-10T09:23:39.158Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/8f/aea9c71cc92bf9b6cc0f7f70df8f0b420636b6c96ef4feee1e16f80f75dd/websockets-16.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0298d07ee155e2e9fda5be8a9042200dd2e3bb0b8a38482156576f863a9d457c", size = 176968, upload-time = "2026-01-10T09:23:41.031Z" },
+ { url = "https://files.pythonhosted.org/packages/9a/3f/f70e03f40ffc9a30d817eef7da1be72ee4956ba8d7255c399a01b135902a/websockets-16.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:a653aea902e0324b52f1613332ddf50b00c06fdaf7e92624fbf8c77c78fa5767", size = 178735, upload-time = "2026-01-10T09:23:42.259Z" },
+ { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" },
+]
+
[[package]]
name = "yarl"
version = "1.22.0"
From 795eb0681aa7dd70b7e18b30203e37712d874ab7 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:16:07 -0800
Subject: [PATCH 02/60] Improve error handling in user verification by
preserving original exception context for invalid tokens
---
osmosis_ai/auth/flow.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/osmosis_ai/auth/flow.py b/osmosis_ai/auth/flow.py
index e65c1353..088a9c44 100644
--- a/osmosis_ai/auth/flow.py
+++ b/osmosis_ai/auth/flow.py
@@ -240,7 +240,7 @@ def _verify_and_get_user_info(
except HTTPError as e:
if e.code == 401:
- raise LoginError("Invalid or expired token") from None
+ raise LoginError("Invalid or expired token") from e
raise LoginError(f"Verification failed: HTTP {e.code}") from e
except URLError as e:
raise LoginError(f"Could not connect to platform: {e.reason}") from e
From 56e3458c0fed08f24c15cd4f69df49bb20aace3f Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:18:04 -0800
Subject: [PATCH 03/60] Add linting and formatting guidelines using Ruff to
README.md
---
README.md | 23 +++++++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/README.md b/README.md
index dc106d27..e406f9e9 100644
--- a/README.md
+++ b/README.md
@@ -397,6 +397,29 @@ Run `python -m pytest` (or any subset under `tests/`) to exercise the updated he
Add additional tests under `tests/` as you extend the library.
+## Linting & Formatting
+
+This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and code formatting. Configuration lives in `pyproject.toml` under `[tool.ruff]`.
+
+```bash
+# Install dev dependencies (includes ruff)
+pip install -e ".[dev]"
+
+# Check for lint errors
+ruff check .
+
+# Auto-fix lint errors where possible
+ruff check --fix .
+
+# Format code
+ruff format .
+
+# Check formatting without modifying files
+ruff format --check .
+```
+
+Please run `ruff check` and `ruff format` before submitting a pull request.
+
## License
MIT License - see [LICENSE](LICENSE) file for details.
From 4aab1c0e0d84815f9e257dcccf4eeb2fae171d0d Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:27:29 -0800
Subject: [PATCH 04/60] Add linting job using Ruff to GitHub Actions workflow
---
.github/workflows/tests.yml | 21 +++++++++++++++++++++
.pre-commit-config.yaml | 7 +++++++
2 files changed, 28 insertions(+)
create mode 100644 .pre-commit-config.yaml
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index a31d7711..b6b032c5 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -7,6 +7,27 @@ on:
pull_request:
jobs:
+ lint:
+ runs-on: ubuntu-latest
+
+ steps:
+ - name: Checkout repository
+ uses: actions/checkout@v4
+
+ - name: Set up Python
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.12"
+
+ - name: Install ruff
+ run: pip install "ruff>=0.9.0,<1.0.0"
+
+ - name: Ruff check
+ run: ruff check .
+
+ - name: Ruff format check
+ run: ruff format --check .
+
pytest:
runs-on: ubuntu-latest
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 00000000..d4de6b03
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,7 @@
+repos:
+ - repo: https://github.com/astral-sh/ruff-pre-commit
+ rev: v0.9.10
+ hooks:
+ - id: ruff
+ args: [--fix]
+ - id: ruff-format
From 5e7a0db3ee4464297d212d1b429f5b170ad248d9 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:28:23 -0800
Subject: [PATCH 05/60] Add pre-commit dependency and update package versions
in pyproject.toml and uv.lock
- Added pre-commit version 4.5.1 to the development dependencies.
- Included new packages cfgv version 3.5.0, distlib version 0.4.0, identify version 2.6.16, and nodeenv version 1.10.0 in the lock file.
- Updated filelock version to 3.24.2 in the lock file.
---
pyproject.toml | 1 +
uv.lock | 75 ++++++++++++++++++++++++++++++++++++++++++++++++--
2 files changed, 73 insertions(+), 3 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 114c66f1..f5abb4f0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,6 +49,7 @@ dev = [
"pytest-asyncio>=0.23.0,<2.0.0",
"anyio>=4.0.0",
"ruff>=0.9.0,<1.0.0",
+ "pre-commit>=4.0.0,<5.0.0",
"fastmcp>=2.0.0",
]
diff --git a/uv.lock b/uv.lock
index 67e1f85d..5e3bca13 100644
--- a/uv.lock
+++ b/uv.lock
@@ -337,6 +337,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ae/3a/dbeec9d1ee0844c679f6bb5d6ad4e9f198b1224f4e7a32825f47f6192b0c/cffi-2.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0a1527a803f0a659de1af2e1fd700213caba79377e27e4693648c2923da066f9", size = 184195, upload-time = "2025-09-08T23:23:43.004Z" },
]
+[[package]]
+name = "cfgv"
+version = "3.5.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4e/b5/721b8799b04bf9afe054a3899c6cf4e880fcf8563cc71c15610242490a0c/cfgv-3.5.0.tar.gz", hash = "sha256:d5b1034354820651caa73ede66a6294d6e95c1b00acc5e9b098e917404669132", size = 7334, upload-time = "2025-11-19T20:55:51.612Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" },
+]
+
[[package]]
name = "charset-normalizer"
version = "3.4.4"
@@ -555,6 +564,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" },
]
+[[package]]
+name = "distlib"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/96/8e/709914eb2b5749865801041647dc7f4e6d00b549cfe88b65ca192995f07c/distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d", size = 614605, upload-time = "2025-07-17T16:52:00.465Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" },
+]
+
[[package]]
name = "distro"
version = "1.9.0"
@@ -744,11 +762,11 @@ wheels = [
[[package]]
name = "filelock"
-version = "3.20.3"
+version = "3.24.2"
source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1d/65/ce7f1b70157833bf3cb851b556a37d4547ceafc158aa9b34b36782f23696/filelock-3.20.3.tar.gz", hash = "sha256:18c57ee915c7ec61cff0ecf7f0f869936c7c30191bb0cf406f1341778d0834e1", size = 19485, upload-time = "2026-01-09T17:55:05.421Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/02/a8/dae62680be63cbb3ff87cfa2f51cf766269514ea5488479d42fec5aa6f3a/filelock-3.24.2.tar.gz", hash = "sha256:c22803117490f156e59fafce621f0550a7a853e2bbf4f87f112b11d469b6c81b", size = 37601, upload-time = "2026-02-16T02:50:45.614Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl", hash = "sha256:4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1", size = 16701, upload-time = "2026-01-09T17:55:04.334Z" },
+ { url = "https://files.pythonhosted.org/packages/e7/04/a94ebfb4eaaa08db56725a40de2887e95de4e8641b9e902c311bfa00aa39/filelock-3.24.2-py3-none-any.whl", hash = "sha256:667d7dc0b7d1e1064dd5f8f8e80bdac157a6482e8d2e02cd16fd3b6b33bd6556", size = 24152, upload-time = "2026-02-16T02:50:44Z" },
]
[[package]]
@@ -1038,6 +1056,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/90/fb/cb8fe5f71d5622427f20bcab9e06a696a5aaf21bfe7bd0a8a0c63c88abf5/huggingface_hub-1.3.1-py3-none-any.whl", hash = "sha256:efbc7f3153cb84e2bb69b62ed90985e21ecc9343d15647a419fc0ee4b85f0ac3", size = 533351, upload-time = "2026-01-09T14:08:14.519Z" },
]
+[[package]]
+name = "identify"
+version = "2.6.16"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/5b/8d/e8b97e6bd3fb6fb271346f7981362f1e04d6a7463abd0de79e1fda17c067/identify-2.6.16.tar.gz", hash = "sha256:846857203b5511bbe94d5a352a48ef2359532bc8f6727b5544077a0dcfb24980", size = 99360, upload-time = "2026-01-12T18:58:58.201Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b8/58/40fbbcefeda82364720eba5cf2270f98496bdfa19ea75b4cccae79c698e6/identify-2.6.16-py2.py3-none-any.whl", hash = "sha256:391ee4d77741d994189522896270b787aed8670389bfd60f326d677d64a6dfb0", size = 99202, upload-time = "2026-01-12T18:58:56.627Z" },
+]
+
[[package]]
name = "idna"
version = "3.11"
@@ -1667,6 +1694,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/da/7d22601b625e241d4f23ef1ebff8acfc60da633c9e7e7922e24d10f592b3/multidict-6.7.0-py3-none-any.whl", hash = "sha256:394fc5c42a333c9ffc3e421a4c85e08580d990e08b99f6bf35b4132114c5dcb3", size = 12317, upload-time = "2025-10-06T14:52:29.272Z" },
]
+[[package]]
+name = "nodeenv"
+version = "1.10.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" },
+]
+
[[package]]
name = "openai"
version = "2.13.0"
@@ -1728,6 +1764,7 @@ dependencies = [
dev = [
{ name = "anyio" },
{ name = "fastmcp" },
+ { name = "pre-commit" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "ruff" },
@@ -1761,6 +1798,7 @@ requires-dist = [
{ name = "fastmcp", marker = "extra == 'mcp'", specifier = ">=2.0.0" },
{ name = "httpx", specifier = ">=0.25.0,<1.0.0" },
{ name = "litellm", specifier = ">=1.40.0,<2.0.0" },
+ { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=4.0.0,<5.0.0" },
{ name = "pyarrow", marker = "extra == 'full'", specifier = ">=14.0.0" },
{ name = "pyarrow", marker = "extra == 'server'", specifier = ">=14.0.0" },
{ name = "pydantic", specifier = ">=2.0.0,<3.0.0" },
@@ -1825,6 +1863,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
]
+[[package]]
+name = "pre-commit"
+version = "4.5.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cfgv" },
+ { name = "identify" },
+ { name = "nodeenv" },
+ { name = "pyyaml" },
+ { name = "virtualenv" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/40/f1/6d86a29246dfd2e9b6237f0b5823717f60cad94d47ddc26afa916d21f525/pre_commit-4.5.1.tar.gz", hash = "sha256:eb545fcff725875197837263e977ea257a402056661f09dae08e4b149b030a61", size = 198232, upload-time = "2025-12-16T21:14:33.552Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl", hash = "sha256:3b3afd891e97337708c1674210f8eba659b52a38ea5f822ff142d10786221f77", size = 226437, upload-time = "2025-12-16T21:14:32.409Z" },
+]
+
[[package]]
name = "prometheus-client"
version = "0.24.1"
@@ -3071,6 +3125,21 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ee/d9/d88e73ca598f4f6ff671fb5fde8a32925c2e08a637303a1d12883c7305fa/uvicorn-0.38.0-py3-none-any.whl", hash = "sha256:48c0afd214ceb59340075b4a052ea1ee91c16fbc2a9b1469cca0e54566977b02", size = 68109, upload-time = "2025-10-18T13:46:42.958Z" },
]
+[[package]]
+name = "virtualenv"
+version = "20.37.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "distlib" },
+ { name = "filelock" },
+ { name = "platformdirs" },
+ { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c1/ef/d9d4ce633df789bf3430bd81fb0d8b9d9465dfc1d1f0deb3fb62cd80f5c2/virtualenv-20.37.0.tar.gz", hash = "sha256:6f7e2064ed470aa7418874e70b6369d53b66bcd9e9fd5389763e96b6c94ccb7c", size = 5864710, upload-time = "2026-02-16T16:17:59.42Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl", hash = "sha256:5d3951c32d57232ae3569d4de4cc256c439e045135ebf43518131175d9be435d", size = 5837480, upload-time = "2026-02-16T16:17:57.341Z" },
+]
+
[[package]]
name = "websockets"
version = "16.0"
From 95634d498c28a2ede3b0a18256510931dc8c078e Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:32:22 -0800
Subject: [PATCH 06/60] Update README.md with pre-commit setup for Ruff and
adjust Ruff version in GitHub Actions workflow
---
.github/workflows/tests.yml | 2 +-
README.md | 8 +++++++-
2 files changed, 8 insertions(+), 2 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index b6b032c5..538a3ad7 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -20,7 +20,7 @@ jobs:
python-version: "3.12"
- name: Install ruff
- run: pip install "ruff>=0.9.0,<1.0.0"
+ run: pip install ruff==0.9.10
- name: Ruff check
run: ruff check .
diff --git a/README.md b/README.md
index e406f9e9..ba14f45f 100644
--- a/README.md
+++ b/README.md
@@ -418,7 +418,13 @@ ruff format .
ruff format --check .
```
-Please run `ruff check` and `ruff format` before submitting a pull request.
+A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
+
+```bash
+pre-commit install
+```
+
+This ensures `ruff check --fix` and `ruff format` run before each commit. Please make sure hooks are installed before submitting a pull request.
## License
From f2032a416bcee3b7d7d8256267323b8a1e51039c Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:34:23 -0800
Subject: [PATCH 07/60] Add non-pep604-isinstance rule to Ruff configuration in
pyproject.toml
---
pyproject.toml | 1 +
1 file changed, 1 insertion(+)
diff --git a/pyproject.toml b/pyproject.toml
index f5abb4f0..280d72ac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -105,6 +105,7 @@ ignore = [
"E402", # module-import-not-at-top-of-file (intentional lazy/conditional imports)
"SIM117", # multiple-with-statements (nested with is often clearer in tests)
"RUF012", # mutable-class-default (common with Pydantic models)
+ "UP038", # non-pep604-isinstance (tuple form is idiomatic and safer)
]
[tool.ruff.lint.isort]
From 11a05344f2005da3b4e7e2ca9e686b6d4b44632f Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:44:05 -0800
Subject: [PATCH 08/60] Update README.md to include instructions for installing
pre-commit hooks for Ruff
---
README.md | 3 +++
1 file changed, 3 insertions(+)
diff --git a/README.md b/README.md
index ba14f45f..6348b453 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,9 @@ pip install -e .
# Install with development dependencies (pytest, formatters, etc.)
pip install -e ".[dev]"
+
+# Install pre-commit hooks (runs ruff on every commit)
+pre-commit install
```
## Quick Start
From b75f7b97a0d8f9be859c92d734ae06a78c52809d Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 11:47:41 -0800
Subject: [PATCH 09/60] remove UP038
---
pyproject.toml | 1 -
1 file changed, 1 deletion(-)
diff --git a/pyproject.toml b/pyproject.toml
index 280d72ac..f5abb4f0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -105,7 +105,6 @@ ignore = [
"E402", # module-import-not-at-top-of-file (intentional lazy/conditional imports)
"SIM117", # multiple-with-statements (nested with is often clearer in tests)
"RUF012", # mutable-class-default (common with Pydantic models)
- "UP038", # non-pep604-isinstance (tuple form is idiomatic and safer)
]
[tool.ruff.lint.isort]
From 34af26d9c67331cc4e4d188240d8b14bf15df06a Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 12:03:04 -0800
Subject: [PATCH 10/60] Update Ruff version to 0.15.1 across configuration
files and documentation for consistency
---
.github/workflows/tests.yml | 2 +-
.pre-commit-config.yaml | 2 +-
README.md | 5 +++++
pyproject.toml | 2 +-
uv.lock | 2 +-
5 files changed, 9 insertions(+), 4 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 538a3ad7..27969c59 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -20,7 +20,7 @@ jobs:
python-version: "3.12"
- name: Install ruff
- run: pip install ruff==0.9.10
+ run: pip install ruff==0.15.1
- name: Ruff check
run: ruff check .
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d4de6b03..d7d2af79 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
- rev: v0.9.10
+ rev: v0.15.1
hooks:
- id: ruff
args: [--fix]
diff --git a/README.md b/README.md
index 6348b453..f1b92936 100644
--- a/README.md
+++ b/README.md
@@ -408,6 +408,9 @@ This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and code form
# Install dev dependencies (includes ruff)
pip install -e ".[dev]"
+# Verify Ruff version
+ruff --version
+
# Check for lint errors
ruff check .
@@ -421,6 +424,8 @@ ruff format .
ruff format --check .
```
+Ruff is pinned to one version across `pyproject.toml`, `.pre-commit-config.yaml`, and CI so local checks, pre-commit hooks, and GitHub Actions produce the same results.
+
A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
```bash
diff --git a/pyproject.toml b/pyproject.toml
index f5abb4f0..dfd5e38b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,7 @@ dev = [
"pytest>=8.0.0,<10.0.0",
"pytest-asyncio>=0.23.0,<2.0.0",
"anyio>=4.0.0",
- "ruff>=0.9.0,<1.0.0",
+ "ruff==0.15.1", # Keep in sync with pre-commit and CI
"pre-commit>=4.0.0,<5.0.0",
"fastmcp>=2.0.0",
]
diff --git a/uv.lock b/uv.lock
index 5e3bca13..9f6b88d5 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1811,7 +1811,7 @@ requires-dist = [
{ name = "requests", specifier = ">=2.0.0,<3.0.0" },
{ name = "rich", marker = "extra == 'full'", specifier = ">=13.0.0" },
{ name = "rich", marker = "extra == 'server'", specifier = ">=13.0.0" },
- { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0,<1.0.0" },
+ { name = "ruff", marker = "extra == 'dev'", specifier = "==0.15.1" },
{ name = "tqdm", specifier = ">=4.0.0,<5.0.0" },
{ name = "uvicorn", marker = "extra == 'full'", specifier = ">=0.23.0,<1.0.0" },
{ name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.23.0,<1.0.0" },
From c6273bac51076c4e0fcba2da40021b85435c4837 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 12:41:05 -0800
Subject: [PATCH 11/60] Remove unused `gather_text_fragments` and
`collect_text_fragments` functions from `shared.py` to clean up the codebase.
---
osmosis_ai/cli_services/shared.py | 140 ------------------------------
1 file changed, 140 deletions(-)
diff --git a/osmosis_ai/cli_services/shared.py b/osmosis_ai/cli_services/shared.py
index 2f6ea730..c1d8c984 100644
--- a/osmosis_ai/cli_services/shared.py
+++ b/osmosis_ai/cli_services/shared.py
@@ -1,6 +1,5 @@
from __future__ import annotations
-from collections.abc import Collection
from statistics import mean, pstdev, pvariance
from typing import Any
@@ -65,142 +64,3 @@ def calculate_stat_deltas(
continue
delta[key] = current_numeric - baseline_value
return delta
-
-
-def gather_text_fragments(
- node: Any,
- fragments: list[str],
- *,
- allow_free_strings: bool = False,
- seen: set[int] | None = None,
- string_key_allowlist: Collection[str] | None = 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) or "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: Collection[str] | None = 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
From 5380954d3342619ec5a5c336b15686fafcad8822 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 12:47:43 -0800
Subject: [PATCH 12/60] address cubic review
---
examples/reward_functions.py | 13 +++++++------
osmosis_ai/utils.py | 17 +++++++++++++----
2 files changed, 20 insertions(+), 10 deletions(-)
diff --git a/examples/reward_functions.py b/examples/reward_functions.py
index 37160381..6a269148 100644
--- a/examples/reward_functions.py
+++ b/examples/reward_functions.py
@@ -2,7 +2,7 @@
Reward function examples using the @osmosis_reward decorator.
This file demonstrates correct and incorrect usage of the decorator,
-which enforces the signature: (solution_str: str, ground_truth: str, extra_info: dict = None) -> float
+which enforces the signature: (solution_str: str, ground_truth: str, extra_info: dict | None = None) -> float
"""
from osmosis_ai import osmosis_reward
@@ -58,10 +58,11 @@ def numeric_tolerance(
return 0.0
-# Only two parameters (extra_info is optional)
@osmosis_reward
-def minimal_reward(solution_str: str, ground_truth: str) -> float:
- """Minimal reward function with just required parameters."""
+def minimal_reward(
+ solution_str: str, ground_truth: str, extra_info: dict | None = None
+) -> float:
+ """Minimal reward function ignoring extra_info."""
return float(solution_str == ground_truth)
@@ -78,11 +79,11 @@ def minimal_reward(solution_str: str, ground_truth: str) -> float:
# return solution_str == ground_truth
# @osmosis_reward # Wrong type annotations
-# def wrong_types(solution_str: int, ground_truth: str, extra_info: dict = None):
+# def wrong_types(solution_str: int, ground_truth: str, extra_info: dict | None = None):
# return str(solution_str) == ground_truth
# @osmosis_reward # Too many parameters
-# def too_many_params(solution_str: str, ground_truth: str, extra_info: dict = None, bonus: float = 0.0):
+# def too_many_params(solution_str: str, ground_truth: str, extra_info: dict | None = None, bonus: float = 0.0):
# return (solution_str == ground_truth) + bonus
# @osmosis_reward # Missing default value for extra_info
diff --git a/osmosis_ai/utils.py b/osmosis_ai/utils.py
index f1295cda..d36b51ff 100644
--- a/osmosis_ai/utils.py
+++ b/osmosis_ai/utils.py
@@ -62,10 +62,19 @@ def osmosis_reward(func: Callable) -> Callable:
raise TypeError(
f"Third parameter must be named 'extra_info', got '{params[2].name}'"
)
- if params[2].annotation is not dict:
- raise TypeError(
- f"Third parameter 'extra_info' must be annotated as dict, got {params[2].annotation}"
- )
+ ann = params[2].annotation
+ if ann is not dict:
+ # Also accept dict | None (PEP 604) and Optional[dict]
+ origin = get_origin(ann)
+ if origin not in _ALLOWED_UNION_ORIGINS:
+ raise TypeError(
+ f"Third parameter 'extra_info' must be annotated as dict or dict | None, got {ann}"
+ )
+ non_none_args = tuple(a for a in get_args(ann) if a is not type(None))
+ if non_none_args != (dict,):
+ raise TypeError(
+ f"Third parameter 'extra_info' must be annotated as dict or dict | None, got {ann}"
+ )
if params[2].default is inspect.Parameter.empty:
raise TypeError(
"Third parameter 'extra_info' must have a default value of None"
From f02395d3b3d336e5e793e2d50a592566ffa4bf9d Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 13:42:19 -0800
Subject: [PATCH 13/60] Add coverage support to testing setup
- Included `pytest-cov` in development dependencies for coverage reporting.
- Configured coverage settings in `pyproject.toml` to specify source and report options.
- Updated GitHub Actions workflow to run tests with coverage.
- Added new tests for deterministic behavior and sequence validation in `fake_token_ids` and `fake_prompt_token_ids` functions.
- Enhanced `RolloutCompletionTracker` tests to verify clearing and recording functionality.
---
.github/workflows/tests.yml | 4 +-
CONTRIBUTING.md | 75 +++
README.md | 67 +-
pyproject.toml | 14 +
tests/test_cli_utils.py | 343 ++++++++++
tests/test_network.py | 597 +++++++++++++++++
tests/test_platform_client.py | 463 +++++++++++++
tests/test_registration.py | 513 +++++++++++++++
tests/test_rollout_server_app.py | 1037 ++++++++++++++++++++++++++++++
tests/test_rollout_testing.py | 471 ++++++++++++++
tests/test_utils_decorators.py | 946 +++++++++++++++++++++++++++
uv.lock | 134 ++++
12 files changed, 4598 insertions(+), 66 deletions(-)
create mode 100644 CONTRIBUTING.md
create mode 100644 tests/test_cli_utils.py
create mode 100644 tests/test_network.py
create mode 100644 tests/test_platform_client.py
create mode 100644 tests/test_registration.py
create mode 100644 tests/test_rollout_server_app.py
create mode 100644 tests/test_utils_decorators.py
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 27969c59..096db083 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -43,5 +43,5 @@ jobs:
- name: Install dependencies
run: pip install -e ".[dev]"
- - name: Run pytest
- run: python -m pytest
+ - name: Run pytest with coverage
+ run: python -m pytest --cov=osmosis_ai --cov-report=term-missing
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 00000000..a8a0ce9c
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,75 @@
+# Contributing
+
+## Setup
+
+```bash
+git clone https://github.com/Osmosis-AI/osmosis-sdk-python
+cd osmosis-sdk-python
+
+# Install with all development dependencies
+pip install -e ".[dev]"
+
+# Or using uv
+uv sync --all-extras
+
+# Install pre-commit hooks
+pre-commit install
+```
+
+## Testing
+
+```bash
+# Run all tests
+pytest tests/
+
+# Run a single test file
+pytest tests/test_rollout_base.py
+
+# Run tests matching a pattern
+pytest -k "test_name"
+
+# Run with coverage report
+pytest --cov=osmosis_ai --cov-report=term-missing
+```
+
+Coverage configuration is in `pyproject.toml` under `[tool.coverage.*]`. CI enforces a minimum coverage threshold.
+
+## Linting & Formatting
+
+This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and code formatting. Configuration lives in `pyproject.toml` under `[tool.ruff]`.
+
+```bash
+# Check for lint errors
+ruff check .
+
+# Auto-fix lint errors where possible
+ruff check --fix .
+
+# Format code
+ruff format .
+
+# Check formatting without modifying files
+ruff format --check .
+```
+
+Ruff is pinned to one version across `pyproject.toml`, `.pre-commit-config.yaml`, and CI so local checks, pre-commit hooks, and GitHub Actions produce the same results.
+
+## Pre-commit Hooks
+
+A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
+
+```bash
+pre-commit install
+```
+
+This ensures `ruff check --fix` and `ruff format` run before each commit. Please make sure hooks are installed before submitting a pull request.
+
+## Pull Requests
+
+1. Fork the repository
+2. Create a feature branch
+3. Make your changes
+4. Run tests and linting
+5. Submit a pull request
+
+CI will run linting (`ruff check` + `ruff format --check`) and tests with coverage on every PR.
diff --git a/README.md b/README.md
index f1b92936..6a4c933e 100644
--- a/README.md
+++ b/README.md
@@ -17,20 +17,7 @@ Requires Python 3.10 or newer.
This installs the Osmosis CLI and pulls in `litellm` (unified LLM interface supporting 100+ providers) along with supporting utilities such as `PyYAML`, `python-dotenv`, and `requests`.
-For development:
-```bash
-git clone https://github.com/Osmosis-AI/osmosis-sdk-python
-cd osmosis-sdk-python
-
-# Install package in editable mode
-pip install -e .
-
-# Install with development dependencies (pytest, formatters, etc.)
-pip install -e ".[dev]"
-
-# Install pre-commit hooks (runs ruff on every commit)
-pre-commit install
-```
+For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
## Quick Start
@@ -390,62 +377,14 @@ PYTHONPATH=. python examples/reward_functions.py
PYTHONPATH=. python examples/rubric_functions.py # Uncomment the provider you need before running
```
-## Testing
-
-Run `python -m pytest` (or any subset under `tests/`) to exercise the updated helpers:
-
-- `tests/test_rubric_eval.py` covers prompt construction for `solution_str` evaluations.
-- `tests/test_cli_services.py` validates dataset parsing, extra-info enrichment, and engine interactions.
-- `tests/test_cli.py` ensures the CLI pathways surface the new fields end to end.
-
-Add additional tests under `tests/` as you extend the library.
-
-## Linting & Formatting
-
-This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and code formatting. Configuration lives in `pyproject.toml` under `[tool.ruff]`.
-
-```bash
-# Install dev dependencies (includes ruff)
-pip install -e ".[dev]"
-
-# Verify Ruff version
-ruff --version
-
-# Check for lint errors
-ruff check .
-
-# Auto-fix lint errors where possible
-ruff check --fix .
-
-# Format code
-ruff format .
-
-# Check formatting without modifying files
-ruff format --check .
-```
-
-Ruff is pinned to one version across `pyproject.toml`, `.pre-commit-config.yaml`, and CI so local checks, pre-commit hooks, and GitHub Actions produce the same results.
-
-A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
-
-```bash
-pre-commit install
-```
+## Contributing
-This ensures `ruff check --fix` and `ruff format` run before each commit. Please make sure hooks are installed before submitting a pull request.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, testing, linting, and PR guidelines.
## License
MIT License - see [LICENSE](LICENSE) file for details.
-## Contributing
-
-1. Fork the repository
-2. Create a feature branch
-3. Make your changes
-4. Run tests and examples
-5. Submit a pull request
-
## Links
- [Homepage](https://github.com/Osmosis-AI/osmosis-sdk-python)
diff --git a/pyproject.toml b/pyproject.toml
index dfd5e38b..74d3fc53 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,6 +47,7 @@ mcp = [
dev = [
"pytest>=8.0.0,<10.0.0",
"pytest-asyncio>=0.23.0,<2.0.0",
+ "pytest-cov>=6.0.0",
"anyio>=4.0.0",
"ruff==0.15.1", # Keep in sync with pre-commit and CI
"pre-commit>=4.0.0,<5.0.0",
@@ -85,6 +86,19 @@ osmosis_ai = ["py.typed"]
[tool.setuptools.dynamic]
version = {attr = "osmosis_ai.consts.PACKAGE_VERSION"}
+[tool.coverage.run]
+source = ["osmosis_ai"]
+omit = ["tests/*", "*/consts.py"]
+
+[tool.coverage.report]
+fail_under = 70
+show_missing = true
+exclude_lines = [
+ "pragma: no cover",
+ "if TYPE_CHECKING:",
+ "if __name__ == .__main__.",
+]
+
[tool.ruff]
target-version = "py310"
line-length = 88
diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py
new file mode 100644
index 00000000..ad10a3e9
--- /dev/null
+++ b/tests/test_cli_utils.py
@@ -0,0 +1,343 @@
+# Copyright 2025 Osmosis AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for osmosis_ai.rollout.cli_utils."""
+
+from __future__ import annotations
+
+import sys
+import textwrap
+
+import pytest
+
+from osmosis_ai.rollout.cli_utils import CLIError, load_agent_loop
+from osmosis_ai.rollout.core.base import RolloutAgentLoop
+
+# =============================================================================
+# CLIError Tests
+# =============================================================================
+
+
+def test_cli_error_is_exception() -> None:
+ """Verify CLIError is a proper Exception subclass."""
+ error = CLIError("something went wrong")
+ assert isinstance(error, Exception)
+ assert str(error) == "something went wrong"
+
+
+def test_cli_error_can_be_raised_and_caught() -> None:
+ """Verify CLIError can be raised and caught specifically."""
+ with pytest.raises(CLIError, match="test message"):
+ raise CLIError("test message")
+
+
+# =============================================================================
+# load_agent_loop - Input Validation
+# =============================================================================
+
+
+def test_load_agent_loop_rejects_missing_colon() -> None:
+ """Verify load_agent_loop raises CLIError when ':' is missing from the module path."""
+ with pytest.raises(CLIError, match="Invalid module path") as exc_info:
+ load_agent_loop("my_module_without_colon")
+ assert "Expected format: 'module.path:attribute_name'" in str(exc_info.value)
+
+
+def test_load_agent_loop_rejects_plain_string() -> None:
+ """Verify load_agent_loop rejects a bare module name with no attribute."""
+ with pytest.raises(CLIError, match="Invalid module path 'agent'"):
+ load_agent_loop("agent")
+
+
+# =============================================================================
+# load_agent_loop - Module Import Errors
+# =============================================================================
+
+
+def test_load_agent_loop_raises_cli_error_for_nonexistent_module() -> None:
+ """Verify load_agent_loop wraps ImportError in CLIError for missing modules."""
+ with pytest.raises(CLIError, match="Cannot import module") as exc_info:
+ load_agent_loop("this_module_definitely_does_not_exist_xyz:some_attr")
+ assert "this_module_definitely_does_not_exist_xyz" in str(exc_info.value)
+
+
+def test_load_agent_loop_import_error_preserves_cause() -> None:
+ """Verify the original ImportError is chained as the cause."""
+ with pytest.raises(CLIError) as exc_info:
+ load_agent_loop("nonexistent_module_abc:attr")
+ assert exc_info.value.__cause__ is not None
+ assert isinstance(exc_info.value.__cause__, ImportError)
+
+
+# =============================================================================
+# load_agent_loop - Attribute Resolution
+# =============================================================================
+
+
+def test_load_agent_loop_raises_cli_error_for_missing_attribute() -> None:
+ """Verify load_agent_loop raises CLIError when the attribute doesn't exist on the module."""
+ # 'os' is always importable but doesn't have 'nonexistent_attr_xyz'
+ with pytest.raises(
+ CLIError, match="has no attribute 'nonexistent_attr_xyz'"
+ ) as exc_info:
+ load_agent_loop("os:nonexistent_attr_xyz")
+ assert "Available attributes:" in str(exc_info.value)
+
+
+def test_load_agent_loop_missing_attribute_lists_available_attrs() -> None:
+ """Verify the error message includes available attributes when attribute is missing."""
+ with pytest.raises(CLIError) as exc_info:
+ load_agent_loop("os:nonexistent_thing")
+ # os module has 'path' as a public attribute
+ assert "path" in str(exc_info.value)
+
+
+# =============================================================================
+# load_agent_loop - Class vs Instance Handling
+# =============================================================================
+
+
+def test_load_agent_loop_accepts_valid_instance(tmp_path) -> None:
+ """Verify load_agent_loop returns a pre-instantiated RolloutAgentLoop instance."""
+ module_code = textwrap.dedent("""\
+ from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext, RolloutResult
+ from osmosis_ai.rollout.core.schemas import RolloutRequest, OpenAIFunctionToolSchema
+
+ class MyAgent(RolloutAgentLoop):
+ name = "test_agent"
+
+ def get_tools(self, request):
+ return []
+
+ async def run(self, ctx):
+ return ctx.complete([])
+
+ agent_instance = MyAgent()
+ """)
+ module_file = tmp_path / "test_agent_instance.py"
+ module_file.write_text(module_code)
+
+ # Temporarily add tmp_path to sys.path
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ # Remove cached module if present
+ sys.modules.pop("test_agent_instance", None)
+ result = load_agent_loop("test_agent_instance:agent_instance")
+ assert isinstance(result, RolloutAgentLoop)
+ assert result.name == "test_agent"
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_agent_instance", None)
+
+
+def test_load_agent_loop_instantiates_valid_class(tmp_path) -> None:
+ """Verify load_agent_loop instantiates a RolloutAgentLoop subclass when given a class."""
+ module_code = textwrap.dedent("""\
+ from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext, RolloutResult
+ from osmosis_ai.rollout.core.schemas import RolloutRequest, OpenAIFunctionToolSchema
+
+ class MyAgent(RolloutAgentLoop):
+ name = "class_agent"
+
+ def get_tools(self, request):
+ return []
+
+ async def run(self, ctx):
+ return ctx.complete([])
+ """)
+ module_file = tmp_path / "test_agent_class.py"
+ module_file.write_text(module_code)
+
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ sys.modules.pop("test_agent_class", None)
+ result = load_agent_loop("test_agent_class:MyAgent")
+ assert isinstance(result, RolloutAgentLoop)
+ assert result.name == "class_agent"
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_agent_class", None)
+
+
+def test_load_agent_loop_rejects_non_rollout_class(tmp_path) -> None:
+ """Verify load_agent_loop raises CLIError for a class that is not a RolloutAgentLoop subclass."""
+ module_code = textwrap.dedent("""\
+ class NotAnAgent:
+ name = "not_agent"
+ """)
+ module_file = tmp_path / "test_non_agent_class.py"
+ module_file.write_text(module_code)
+
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ sys.modules.pop("test_non_agent_class", None)
+ with pytest.raises(CLIError, match="not a RolloutAgentLoop subclass"):
+ load_agent_loop("test_non_agent_class:NotAnAgent")
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_non_agent_class", None)
+
+
+def test_load_agent_loop_rejects_non_rollout_instance(tmp_path) -> None:
+ """Verify load_agent_loop raises CLIError for an object that is not a RolloutAgentLoop instance."""
+ module_code = textwrap.dedent("""\
+ class NotAnAgent:
+ pass
+
+ not_agent_instance = NotAnAgent()
+ """)
+ module_file = tmp_path / "test_non_agent_inst.py"
+ module_file.write_text(module_code)
+
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ sys.modules.pop("test_non_agent_inst", None)
+ with pytest.raises(
+ CLIError, match="must be a RolloutAgentLoop instance or subclass"
+ ):
+ load_agent_loop("test_non_agent_inst:not_agent_instance")
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_non_agent_inst", None)
+
+
+def test_load_agent_loop_rejects_plain_value(tmp_path) -> None:
+ """Verify load_agent_loop raises CLIError when attribute is a plain value (e.g. string, int)."""
+ module_code = textwrap.dedent("""\
+ some_string = "not an agent"
+ """)
+ module_file = tmp_path / "test_plain_value.py"
+ module_file.write_text(module_code)
+
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ sys.modules.pop("test_plain_value", None)
+ with pytest.raises(
+ CLIError, match="must be a RolloutAgentLoop instance or subclass"
+ ):
+ load_agent_loop("test_plain_value:some_string")
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_plain_value", None)
+
+
+def test_load_agent_loop_class_instantiation_failure(tmp_path) -> None:
+ """Verify load_agent_loop raises CLIError when class constructor raises an exception."""
+ module_code = textwrap.dedent("""\
+ from osmosis_ai.rollout.core.base import RolloutAgentLoop, RolloutContext, RolloutResult
+ from osmosis_ai.rollout.core.schemas import RolloutRequest, OpenAIFunctionToolSchema
+
+ class BrokenAgent(RolloutAgentLoop):
+ name = "broken_agent"
+
+ def __init__(self):
+ raise RuntimeError("constructor exploded")
+
+ def get_tools(self, request):
+ return []
+
+ async def run(self, ctx):
+ return ctx.complete([])
+ """)
+ module_file = tmp_path / "test_broken_agent.py"
+ module_file.write_text(module_code)
+
+ original_path = sys.path.copy()
+ sys.path.insert(0, str(tmp_path))
+ try:
+ sys.modules.pop("test_broken_agent", None)
+ with pytest.raises(
+ CLIError, match="Cannot instantiate 'BrokenAgent'"
+ ) as exc_info:
+ load_agent_loop("test_broken_agent:BrokenAgent")
+ assert "constructor exploded" in str(exc_info.value)
+ finally:
+ sys.path[:] = original_path
+ sys.modules.pop("test_broken_agent", None)
+
+
+# =============================================================================
+# load_agent_loop - sys.path Modification
+# =============================================================================
+
+
+def test_load_agent_loop_adds_cwd_to_sys_path_if_missing() -> None:
+ """Verify load_agent_loop adds cwd to sys.path when it's not already present."""
+ import os
+
+ cwd = os.getcwd()
+
+ # Temporarily remove cwd from sys.path
+ original_path = sys.path.copy()
+ sys.path = [p for p in sys.path if p != cwd]
+
+ try:
+ assert cwd not in sys.path
+ # This will fail to import but that's fine - we just want to check sys.path
+ with pytest.raises(CLIError):
+ load_agent_loop("nonexistent_xyz_module:attr")
+ assert cwd in sys.path
+ finally:
+ sys.path[:] = original_path
+
+
+def test_load_agent_loop_does_not_duplicate_cwd_in_sys_path() -> None:
+ """Verify load_agent_loop does not add cwd to sys.path if it's already there."""
+ import os
+
+ cwd = os.getcwd()
+
+ # Ensure cwd IS in sys.path
+ if cwd not in sys.path:
+ sys.path.insert(0, cwd)
+
+ count_before = sys.path.count(cwd)
+
+ try:
+ with pytest.raises(CLIError):
+ load_agent_loop("nonexistent_xyz_module:attr")
+ count_after = sys.path.count(cwd)
+ assert count_after == count_before
+ finally:
+ pass
+
+
+# =============================================================================
+# load_agent_loop - Module Path Parsing
+# =============================================================================
+
+
+def test_load_agent_loop_rsplit_handles_multiple_colons() -> None:
+ """Verify load_agent_loop uses rsplit so only the last colon is treated as separator."""
+ # A path like "package.module:submodule:attr" should split as
+ # module_name="package.module:submodule", attr_name="attr"
+ # This will fail at import (no such module) but demonstrates rsplit behavior.
+ with pytest.raises(
+ CLIError, match=r"Cannot import module 'package\.module:submodule'"
+ ):
+ load_agent_loop("package.module:submodule:attr")
+
+
+def test_load_agent_loop_with_dotted_module_path() -> None:
+ """Verify load_agent_loop handles dotted module paths correctly."""
+ # os.path is a valid module; 'join' is a valid attribute (but not an agent)
+ with pytest.raises(
+ CLIError, match="must be a RolloutAgentLoop instance or subclass"
+ ):
+ load_agent_loop("os.path:join")
diff --git a/tests/test_network.py b/tests/test_network.py
new file mode 100644
index 00000000..4e6b48e2
--- /dev/null
+++ b/tests/test_network.py
@@ -0,0 +1,597 @@
+"""Tests for osmosis_ai.rollout.network module.
+
+Covers IP validation helpers (pure logic), cloud metadata detection (mocked HTTP),
+external service fallback, and the main detect_public_ip orchestrator.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import requests
+
+from osmosis_ai.rollout.network import (
+ PublicIPDetectionError,
+ _get_aws_public_ip,
+ _get_azure_public_ip,
+ _get_gcp_public_ip,
+ detect_from_cloud_metadata,
+ detect_from_external_services,
+ detect_public_ip,
+ is_private_ip,
+ is_valid_hostname_or_ip,
+ validate_ipv4,
+)
+
+# =============================================================================
+# validate_ipv4 Tests
+# =============================================================================
+
+
+class TestValidateIPv4:
+ """Tests for the validate_ipv4 helper."""
+
+ def test_standard_public_ip(self) -> None:
+ """Valid public IPv4 addresses should be accepted."""
+ assert validate_ipv4("8.8.8.8") is True
+ assert validate_ipv4("54.123.45.67") is True
+ assert validate_ipv4("1.2.3.4") is True
+
+ def test_private_ip_ranges(self) -> None:
+ """Private IPv4 addresses are still valid IPv4."""
+ assert validate_ipv4("192.168.1.1") is True
+ assert validate_ipv4("10.0.0.1") is True
+ assert validate_ipv4("172.16.0.1") is True
+
+ def test_boundary_values(self) -> None:
+ """Boundary addresses (0.0.0.0, 255.255.255.255) are valid IPv4."""
+ assert validate_ipv4("0.0.0.0") is True
+ assert validate_ipv4("255.255.255.255") is True
+
+ def test_loopback(self) -> None:
+ """Loopback address is valid IPv4."""
+ assert validate_ipv4("127.0.0.1") is True
+
+ def test_octet_out_of_range(self) -> None:
+ """Octets > 255 should be rejected."""
+ assert validate_ipv4("256.1.1.1") is False
+ assert validate_ipv4("1.256.1.1") is False
+ assert validate_ipv4("1.1.1.256") is False
+
+ def test_not_an_ip_at_all(self) -> None:
+ """Non-IP strings should be rejected."""
+ assert validate_ipv4("not-an-ip") is False
+ assert validate_ipv4("hello world") is False
+ assert validate_ipv4("") is False
+
+ def test_ipv6_is_rejected(self) -> None:
+ """IPv6 addresses should be rejected (this function is IPv4-only)."""
+ assert validate_ipv4("::1") is False
+ assert validate_ipv4("2001:db8::1") is False
+ assert validate_ipv4("fe80::1") is False
+
+ def test_too_few_or_too_many_octets(self) -> None:
+ """Addresses with wrong number of octets should be rejected."""
+ assert validate_ipv4("1.2.3") is False
+ assert validate_ipv4("1.2.3.4.5") is False
+
+ def test_leading_zeros(self) -> None:
+ """Leading zeros are rejected by Python's ipaddress module."""
+ assert validate_ipv4("01.02.03.04") is False
+
+ def test_whitespace_not_stripped(self) -> None:
+ """Addresses with surrounding whitespace should be rejected."""
+ assert validate_ipv4(" 1.2.3.4") is False
+ assert validate_ipv4("1.2.3.4 ") is False
+
+
+# =============================================================================
+# is_private_ip Tests
+# =============================================================================
+
+
+class TestIsPrivateIP:
+ """Tests for the is_private_ip helper."""
+
+ def test_rfc1918_class_a(self) -> None:
+ """10.0.0.0/8 range should be private."""
+ assert is_private_ip("10.0.0.1") is True
+ assert is_private_ip("10.255.255.255") is True
+
+ def test_rfc1918_class_b(self) -> None:
+ """172.16.0.0/12 range should be private."""
+ assert is_private_ip("172.16.0.1") is True
+ assert is_private_ip("172.31.255.255") is True
+
+ def test_rfc1918_class_c(self) -> None:
+ """192.168.0.0/16 range should be private."""
+ assert is_private_ip("192.168.0.1") is True
+ assert is_private_ip("192.168.255.255") is True
+
+ def test_loopback_is_private(self) -> None:
+ """Loopback addresses are considered private by Python ipaddress."""
+ assert is_private_ip("127.0.0.1") is True
+
+ def test_public_ips(self) -> None:
+ """Public IPs should not be private."""
+ assert is_private_ip("54.123.45.67") is False
+ assert is_private_ip("8.8.8.8") is False
+ assert is_private_ip("1.1.1.1") is False
+
+ def test_invalid_input_returns_false(self) -> None:
+ """Invalid IP strings should return False, not raise."""
+ assert is_private_ip("not-an-ip") is False
+ assert is_private_ip("") is False
+ assert is_private_ip("999.999.999.999") is False
+
+
+# =============================================================================
+# is_valid_hostname_or_ip Tests
+# =============================================================================
+
+
+class TestIsValidHostnameOrIP:
+ """Tests for the is_valid_hostname_or_ip helper."""
+
+ def test_valid_ipv4_addresses(self) -> None:
+ """IPv4 addresses should be accepted."""
+ assert is_valid_hostname_or_ip("192.168.1.1") is True
+ assert is_valid_hostname_or_ip("8.8.8.8") is True
+ assert is_valid_hostname_or_ip("0.0.0.0") is True
+
+ def test_valid_hostnames(self) -> None:
+ """Standard hostnames should be accepted."""
+ assert is_valid_hostname_or_ip("example.com") is True
+ assert is_valid_hostname_or_ip("sub.example.com") is True
+ assert is_valid_hostname_or_ip("localhost") is True
+
+ def test_single_character_hostname(self) -> None:
+ """Single character hostnames are valid per the regex pattern."""
+ assert is_valid_hostname_or_ip("a") is True
+
+ def test_hostname_with_hyphens(self) -> None:
+ """Hostnames with internal hyphens should be accepted."""
+ assert is_valid_hostname_or_ip("my-host") is True
+ assert is_valid_hostname_or_ip("my-long-hostname.example.com") is True
+
+ def test_hostname_with_underscores(self) -> None:
+ """Hostnames with underscores should be accepted (practical tolerance)."""
+ assert is_valid_hostname_or_ip("my_host") is True
+ assert is_valid_hostname_or_ip("my_host.example.com") is True
+
+ def test_host_with_port(self) -> None:
+ """host:port format should be accepted."""
+ assert is_valid_hostname_or_ip("example.com:8080") is True
+ assert is_valid_hostname_or_ip("192.168.1.1:8080") is True
+ assert is_valid_hostname_or_ip("localhost:3000") is True
+
+ def test_empty_string_rejected(self) -> None:
+ """Empty string should be rejected."""
+ assert is_valid_hostname_or_ip("") is False
+
+ def test_string_with_spaces_rejected(self) -> None:
+ """Strings containing spaces should be rejected."""
+ assert is_valid_hostname_or_ip("has spaces") is False
+ assert is_valid_hostname_or_ip("has space.com") is False
+
+ def test_too_long_rejected(self) -> None:
+ """Values exceeding max length should be rejected."""
+ # Total length > 260 chars
+ long_value = "a" * 261
+ assert is_valid_hostname_or_ip(long_value) is False
+
+ def test_hostname_label_too_long(self) -> None:
+ """Hostname total > 253 chars (without port) should be rejected."""
+ # Host part > 253 chars
+ long_host = "a" * 254
+ assert is_valid_hostname_or_ip(long_host) is False
+
+ def test_colon_with_non_numeric_port_treated_as_hostname(self) -> None:
+ """A colon followed by non-digits is not treated as host:port splitting."""
+ # "host:abc" - the port is not numeric, so the full string is checked as hostname
+ # This should fail because ':' is not in the allowed hostname chars
+ assert is_valid_hostname_or_ip("host:abc") is False
+
+ def test_numeric_hostname(self) -> None:
+ """Purely numeric strings that are not valid IPs should still be checked."""
+ # "1" is a valid single-char hostname
+ assert is_valid_hostname_or_ip("1") is True
+
+
+# =============================================================================
+# AWS Metadata Detection Tests
+# =============================================================================
+
+
+class TestGetAWSPublicIP:
+ """Tests for _get_aws_public_ip with mocked HTTP."""
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_successful_detection(self, mock_requests: MagicMock) -> None:
+ """AWS IMDSv2 two-step flow returns a valid public IP."""
+ mock_token_resp = MagicMock()
+ mock_token_resp.text = "test-token-abc"
+ mock_token_resp.raise_for_status = MagicMock()
+
+ mock_ip_resp = MagicMock()
+ mock_ip_resp.status_code = 200
+ mock_ip_resp.text = "54.123.45.67"
+ mock_ip_resp.raise_for_status = MagicMock()
+
+ mock_requests.put.return_value = mock_token_resp
+ mock_requests.get.return_value = mock_ip_resp
+
+ result = _get_aws_public_ip()
+
+ assert result == "54.123.45.67"
+ # Verify token request was made with PUT
+ mock_requests.put.assert_called_once()
+ put_args = mock_requests.put.call_args
+ assert "169.254.169.254" in put_args[0][0]
+ assert "X-aws-ec2-metadata-token-ttl-seconds" in put_args[1]["headers"]
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_no_public_ip_returns_none(self, mock_requests: MagicMock) -> None:
+ """AWS returns 404 when instance has no public IP (private subnet)."""
+ mock_token_resp = MagicMock()
+ mock_token_resp.text = "test-token"
+ mock_token_resp.raise_for_status = MagicMock()
+
+ mock_ip_resp = MagicMock()
+ mock_ip_resp.status_code = 404
+
+ mock_requests.put.return_value = mock_token_resp
+ mock_requests.get.return_value = mock_ip_resp
+
+ result = _get_aws_public_ip()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_token_request_fails(self, mock_requests: MagicMock) -> None:
+ """When token request fails, returns None gracefully."""
+ mock_requests.put.side_effect = requests.ConnectionError("timeout")
+
+ result = _get_aws_public_ip()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_invalid_ip_returned(self, mock_requests: MagicMock) -> None:
+ """When metadata returns non-IP text, returns None."""
+ mock_token_resp = MagicMock()
+ mock_token_resp.text = "token"
+ mock_token_resp.raise_for_status = MagicMock()
+
+ mock_ip_resp = MagicMock()
+ mock_ip_resp.status_code = 200
+ mock_ip_resp.text = "not-a-valid-ip"
+ mock_ip_resp.raise_for_status = MagicMock()
+
+ mock_requests.put.return_value = mock_token_resp
+ mock_requests.get.return_value = mock_ip_resp
+
+ result = _get_aws_public_ip()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_empty_ip_response(self, mock_requests: MagicMock) -> None:
+ """When metadata returns empty text, returns None."""
+ mock_token_resp = MagicMock()
+ mock_token_resp.text = "token"
+ mock_token_resp.raise_for_status = MagicMock()
+
+ mock_ip_resp = MagicMock()
+ mock_ip_resp.status_code = 200
+ mock_ip_resp.text = " "
+ mock_ip_resp.raise_for_status = MagicMock()
+
+ mock_requests.put.return_value = mock_token_resp
+ mock_requests.get.return_value = mock_ip_resp
+
+ result = _get_aws_public_ip()
+ assert result is None
+
+
+# =============================================================================
+# GCP Metadata Detection Tests
+# =============================================================================
+
+
+class TestGetGCPPublicIP:
+ """Tests for _get_gcp_public_ip with mocked HTTP."""
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_successful_detection(self, mock_requests: MagicMock) -> None:
+ """GCP metadata returns a valid public IP."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 200
+ mock_resp.text = "35.200.100.50"
+ mock_resp.raise_for_status = MagicMock()
+
+ mock_requests.get.return_value = mock_resp
+
+ result = _get_gcp_public_ip()
+
+ assert result == "35.200.100.50"
+ get_args = mock_requests.get.call_args
+ assert "Metadata-Flavor" in get_args[1]["headers"]
+ assert get_args[1]["headers"]["Metadata-Flavor"] == "Google"
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_no_external_ip_returns_none(self, mock_requests: MagicMock) -> None:
+ """GCP returns 404 when instance has no external IP."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 404
+
+ mock_requests.get.return_value = mock_resp
+
+ result = _get_gcp_public_ip()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_connection_error_returns_none(self, mock_requests: MagicMock) -> None:
+ """When GCP metadata service is unreachable, returns None."""
+ mock_requests.get.side_effect = requests.ConnectionError("timeout")
+
+ result = _get_gcp_public_ip()
+ assert result is None
+
+
+# =============================================================================
+# Azure Metadata Detection Tests
+# =============================================================================
+
+
+class TestGetAzurePublicIP:
+ """Tests for _get_azure_public_ip with mocked HTTP."""
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_successful_detection(self, mock_requests: MagicMock) -> None:
+ """Azure IMDS returns a valid public IP."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 200
+ mock_resp.text = "20.50.100.200"
+ mock_resp.raise_for_status = MagicMock()
+
+ mock_requests.get.return_value = mock_resp
+
+ result = _get_azure_public_ip()
+
+ assert result == "20.50.100.200"
+ get_args = mock_requests.get.call_args
+ assert get_args[1]["headers"]["Metadata"] == "true"
+ assert "api-version" in get_args[1]["params"]
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_no_public_ip_returns_none(self, mock_requests: MagicMock) -> None:
+ """Azure returns 404 when instance has no public IP."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 404
+
+ mock_requests.get.return_value = mock_resp
+
+ result = _get_azure_public_ip()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_connection_error_returns_none(self, mock_requests: MagicMock) -> None:
+ """When Azure IMDS is unreachable, returns None."""
+ mock_requests.get.side_effect = requests.Timeout("timeout")
+
+ result = _get_azure_public_ip()
+ assert result is None
+
+
+# =============================================================================
+# detect_from_cloud_metadata Tests
+# =============================================================================
+
+
+class TestDetectFromCloudMetadata:
+ """Tests for the parallel cloud metadata detection orchestrator."""
+
+ @patch("osmosis_ai.rollout.network._get_azure_public_ip", return_value=None)
+ @patch("osmosis_ai.rollout.network._get_gcp_public_ip", return_value=None)
+ @patch(
+ "osmosis_ai.rollout.network._get_aws_public_ip",
+ return_value="54.123.45.67",
+ )
+ def test_returns_first_successful_result(
+ self, mock_aws: MagicMock, mock_gcp: MagicMock, mock_azure: MagicMock
+ ) -> None:
+ """When AWS succeeds, returns its IP even if others fail."""
+ result = detect_from_cloud_metadata()
+ assert result == "54.123.45.67"
+
+ @patch("osmosis_ai.rollout.network._get_azure_public_ip", return_value=None)
+ @patch(
+ "osmosis_ai.rollout.network._get_gcp_public_ip",
+ return_value="35.200.100.50",
+ )
+ @patch("osmosis_ai.rollout.network._get_aws_public_ip", return_value=None)
+ def test_returns_gcp_when_aws_fails(
+ self, mock_aws: MagicMock, mock_gcp: MagicMock, mock_azure: MagicMock
+ ) -> None:
+ """When AWS returns None but GCP succeeds, returns GCP's IP."""
+ result = detect_from_cloud_metadata()
+ assert result == "35.200.100.50"
+
+ @patch("osmosis_ai.rollout.network._get_azure_public_ip", return_value=None)
+ @patch("osmosis_ai.rollout.network._get_gcp_public_ip", return_value=None)
+ @patch("osmosis_ai.rollout.network._get_aws_public_ip", return_value=None)
+ def test_returns_none_when_all_fail(
+ self, mock_aws: MagicMock, mock_gcp: MagicMock, mock_azure: MagicMock
+ ) -> None:
+ """When all cloud providers return None, returns None."""
+ result = detect_from_cloud_metadata()
+ assert result is None
+
+ @patch(
+ "osmosis_ai.rollout.network._get_azure_public_ip",
+ side_effect=Exception("crash"),
+ )
+ @patch(
+ "osmosis_ai.rollout.network._get_gcp_public_ip",
+ side_effect=Exception("crash"),
+ )
+ @patch(
+ "osmosis_ai.rollout.network._get_aws_public_ip",
+ side_effect=Exception("crash"),
+ )
+ def test_handles_exceptions_from_all_providers(
+ self, mock_aws: MagicMock, mock_gcp: MagicMock, mock_azure: MagicMock
+ ) -> None:
+ """When all providers raise exceptions, returns None without crashing."""
+ result = detect_from_cloud_metadata()
+ assert result is None
+
+
+# =============================================================================
+# detect_from_external_services Tests
+# =============================================================================
+
+
+class TestDetectFromExternalServices:
+ """Tests for external IP service fallback detection."""
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_returns_first_valid_ip(self, mock_requests: MagicMock) -> None:
+ """When an external service returns a valid IP, it is returned."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 200
+ mock_resp.text = "203.0.113.42\n"
+ mock_resp.raise_for_status = MagicMock()
+
+ mock_requests.get.return_value = mock_resp
+
+ result = detect_from_external_services()
+ assert result == "203.0.113.42"
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_returns_none_when_all_services_fail(
+ self, mock_requests: MagicMock
+ ) -> None:
+ """When all external services raise errors, returns None."""
+ mock_requests.get.side_effect = requests.ConnectionError("unreachable")
+
+ result = detect_from_external_services()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_returns_none_when_services_return_invalid_ip(
+ self, mock_requests: MagicMock
+ ) -> None:
+ """When services return non-IP text, returns None."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 200
+ mock_resp.text = "error page"
+ mock_resp.raise_for_status = MagicMock()
+
+ mock_requests.get.return_value = mock_resp
+
+ result = detect_from_external_services()
+ assert result is None
+
+ @patch("osmosis_ai.rollout.network.requests")
+ def test_returns_none_when_services_return_empty(
+ self, mock_requests: MagicMock
+ ) -> None:
+ """When services return empty text, returns None."""
+ mock_resp = MagicMock()
+ mock_resp.status_code = 200
+ mock_resp.text = ""
+ mock_resp.raise_for_status = MagicMock()
+
+ mock_requests.get.return_value = mock_resp
+
+ result = detect_from_external_services()
+ assert result is None
+
+
+# =============================================================================
+# detect_public_ip (Main Orchestrator) Tests
+# =============================================================================
+
+
+class TestDetectPublicIP:
+ """Tests for the main detect_public_ip orchestration function."""
+
+ @patch("osmosis_ai.rollout.network.detect_from_external_services")
+ @patch("osmosis_ai.rollout.network.detect_from_cloud_metadata")
+ def test_prefers_cloud_metadata_over_external(
+ self, mock_cloud: MagicMock, mock_external: MagicMock
+ ) -> None:
+ """Cloud metadata result is returned without calling external services."""
+ mock_cloud.return_value = "54.123.45.67"
+
+ result = detect_public_ip()
+
+ assert result == "54.123.45.67"
+ mock_cloud.assert_called_once()
+ mock_external.assert_not_called()
+
+ @patch("osmosis_ai.rollout.network.detect_from_external_services")
+ @patch("osmosis_ai.rollout.network.detect_from_cloud_metadata")
+ def test_falls_back_to_external_when_cloud_fails(
+ self, mock_cloud: MagicMock, mock_external: MagicMock
+ ) -> None:
+ """When cloud metadata returns None, falls back to external services."""
+ mock_cloud.return_value = None
+ mock_external.return_value = "203.0.113.42"
+
+ result = detect_public_ip()
+
+ assert result == "203.0.113.42"
+ mock_cloud.assert_called_once()
+ mock_external.assert_called_once()
+
+ @patch("osmosis_ai.rollout.network.detect_from_external_services")
+ @patch("osmosis_ai.rollout.network.detect_from_cloud_metadata")
+ def test_raises_error_when_all_methods_fail(
+ self, mock_cloud: MagicMock, mock_external: MagicMock
+ ) -> None:
+ """When both cloud and external fail, raises PublicIPDetectionError."""
+ mock_cloud.return_value = None
+ mock_external.return_value = None
+
+ with pytest.raises(PublicIPDetectionError, match="Failed to detect public IP"):
+ detect_public_ip()
+
+ @patch("osmosis_ai.rollout.network.detect_from_external_services")
+ @patch("osmosis_ai.rollout.network.detect_from_cloud_metadata")
+ def test_error_message_is_informative(
+ self, mock_cloud: MagicMock, mock_external: MagicMock
+ ) -> None:
+ """The error message should mention all detection methods that were tried."""
+ mock_cloud.return_value = None
+ mock_external.return_value = None
+
+ with pytest.raises(PublicIPDetectionError) as exc_info:
+ detect_public_ip()
+
+ error_message = str(exc_info.value)
+ assert "Cloud metadata" in error_message
+ assert "External IP services" in error_message
+
+
+# =============================================================================
+# PublicIPDetectionError Tests
+# =============================================================================
+
+
+class TestPublicIPDetectionError:
+ """Tests for the custom exception class."""
+
+ def test_is_exception_subclass(self) -> None:
+ """PublicIPDetectionError should be a proper Exception subclass."""
+ assert issubclass(PublicIPDetectionError, Exception)
+
+ def test_can_be_instantiated_with_message(self) -> None:
+ """Exception can carry a custom message."""
+ err = PublicIPDetectionError("custom message")
+ assert str(err) == "custom message"
+
+ def test_can_be_caught_as_exception(self) -> None:
+ """PublicIPDetectionError can be caught via pytest.raises."""
+ with pytest.raises(PublicIPDetectionError):
+ raise PublicIPDetectionError("test")
diff --git a/tests/test_platform_client.py b/tests/test_platform_client.py
new file mode 100644
index 00000000..0ada2314
--- /dev/null
+++ b/tests/test_platform_client.py
@@ -0,0 +1,463 @@
+"""Tests for osmosis_ai.auth.platform_client."""
+
+from __future__ import annotations
+
+import http.client
+import json
+from datetime import datetime, timedelta, timezone
+from io import BytesIO
+from typing import Any
+from unittest.mock import MagicMock, patch
+from urllib.error import HTTPError, URLError
+
+import pytest
+
+from osmosis_ai.auth.credentials import OrganizationInfo, UserInfo, WorkspaceCredentials
+from osmosis_ai.auth.platform_client import (
+ AuthenticationExpiredError,
+ PlatformAPIError,
+ _handle_401_and_cleanup,
+ platform_request,
+)
+
+# =============================================================================
+# Helper: Create real WorkspaceCredentials for testing
+# =============================================================================
+
+
+def _make_credentials(
+ access_token: str = "test-token-abc123",
+ org_name: str = "TestOrg",
+) -> WorkspaceCredentials:
+ """Create valid WorkspaceCredentials for testing."""
+ now = datetime.now(timezone.utc)
+ return WorkspaceCredentials(
+ access_token=access_token,
+ token_type="Bearer",
+ expires_at=now + timedelta(days=30),
+ user=UserInfo(id="user_1", email="user@example.com", name="Test User"),
+ organization=OrganizationInfo(id="org_1", name=org_name, role="member"),
+ created_at=now,
+ )
+
+
+def _make_http_response(data: dict[str, Any]) -> MagicMock:
+ """Create a mock HTTP response that behaves like urlopen() return value."""
+ body = json.dumps(data).encode()
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+ return mock_resp
+
+
+def _make_http_error(
+ code: int, body: str = "", url: str = "https://platform.osmosis.ai/api/test"
+) -> HTTPError:
+ """Create an HTTPError with readable body."""
+ fp = BytesIO(body.encode("utf-8")) if body else BytesIO(b"")
+ return HTTPError(
+ url=url, code=code, msg=http.client.responses.get(code, "Error"), hdrs={}, fp=fp
+ )
+
+
+# =============================================================================
+# Exception Classes Tests
+# =============================================================================
+
+
+class TestExceptionClasses:
+ """Tests for AuthenticationExpiredError and PlatformAPIError."""
+
+ def test_authentication_expired_error_is_exception(self) -> None:
+ """Verify AuthenticationExpiredError is a proper Exception subclass."""
+ err = AuthenticationExpiredError("session expired")
+ assert isinstance(err, Exception)
+ assert str(err) == "session expired"
+
+ def test_platform_api_error_without_status_code(self) -> None:
+ """Verify PlatformAPIError can be created without a status code."""
+ err = PlatformAPIError("connection failed")
+ assert str(err) == "connection failed"
+ assert err.status_code is None
+
+ def test_platform_api_error_with_status_code(self) -> None:
+ """Verify PlatformAPIError stores the HTTP status code."""
+ err = PlatformAPIError("not found", status_code=404)
+ assert err.status_code == 404
+ assert "not found" in str(err)
+
+
+# =============================================================================
+# _handle_401_and_cleanup Tests
+# =============================================================================
+
+
+class TestHandle401AndCleanup:
+ """Tests for the _handle_401_and_cleanup function."""
+
+ @patch("osmosis_ai.auth.platform_client.delete_workspace_credentials")
+ @patch("osmosis_ai.auth.platform_client.get_active_workspace")
+ def test_deletes_active_workspace_credentials(
+ self, mock_get_ws: MagicMock, mock_delete: MagicMock
+ ) -> None:
+ """Verify 401 handler deletes credentials for the active workspace."""
+ mock_get_ws.return_value = "MyWorkspace"
+
+ with pytest.raises(AuthenticationExpiredError, match="expired or been revoked"):
+ _handle_401_and_cleanup()
+
+ mock_delete.assert_called_once_with("MyWorkspace")
+
+ @patch("osmosis_ai.auth.platform_client.delete_workspace_credentials")
+ @patch("osmosis_ai.auth.platform_client.get_active_workspace")
+ def test_no_active_workspace_skips_delete(
+ self, mock_get_ws: MagicMock, mock_delete: MagicMock
+ ) -> None:
+ """Verify 401 handler does not call delete when no workspace is active."""
+ mock_get_ws.return_value = None
+
+ with pytest.raises(AuthenticationExpiredError, match="osmosis login"):
+ _handle_401_and_cleanup()
+
+ mock_delete.assert_not_called()
+
+ @patch("osmosis_ai.auth.platform_client.delete_workspace_credentials")
+ @patch("osmosis_ai.auth.platform_client.get_active_workspace")
+ def test_always_raises_even_after_cleanup(
+ self, mock_get_ws: MagicMock, mock_delete: MagicMock
+ ) -> None:
+ """Verify 401 handler always raises AuthenticationExpiredError."""
+ mock_get_ws.return_value = "SomeWorkspace"
+ mock_delete.return_value = True
+
+ with pytest.raises(AuthenticationExpiredError):
+ _handle_401_and_cleanup()
+
+
+# =============================================================================
+# platform_request Tests
+# =============================================================================
+
+
+class TestPlatformRequest:
+ """Tests for the platform_request function."""
+
+ # -------------------------------------------------------------------------
+ # Credential Loading
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.load_credentials")
+ def test_raises_when_no_credentials_found(self, mock_load: MagicMock) -> None:
+ """Verify AuthenticationExpiredError when no credentials exist."""
+ mock_load.return_value = None
+
+ with pytest.raises(AuthenticationExpiredError, match="No valid credentials"):
+ platform_request("/api/test")
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ @patch("osmosis_ai.auth.platform_client.load_credentials")
+ def test_uses_loaded_credentials_when_none_provided(
+ self, mock_load: MagicMock, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify credentials are loaded from storage when not explicitly provided."""
+ creds = _make_credentials(access_token="loaded-token")
+ mock_load.return_value = creds
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+
+ platform_request("/api/test")
+
+ mock_load.assert_called_once()
+ # Verify the loaded token was used in the request
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.get_header("Authorization") == "Bearer loaded-token"
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ @patch("osmosis_ai.auth.platform_client.load_credentials")
+ def test_uses_explicit_credentials_when_provided(
+ self, mock_load: MagicMock, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify explicit credentials bypass loading from storage."""
+ explicit_creds = _make_credentials(access_token="explicit-token")
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+
+ platform_request("/api/test", credentials=explicit_creds)
+
+ mock_load.assert_not_called()
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.get_header("Authorization") == "Bearer explicit-token"
+
+ # -------------------------------------------------------------------------
+ # Request Construction
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_constructs_correct_url(self, mock_urlopen: MagicMock) -> None:
+ """Verify the full URL is built from PLATFORM_URL + endpoint."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials()
+
+ with patch(
+ "osmosis_ai.auth.platform_client.PLATFORM_URL", "https://test.osmosis.ai"
+ ):
+ platform_request("/api/v1/verify", credentials=creds)
+
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.full_url == "https://test.osmosis.ai/api/v1/verify"
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_sets_required_headers(self, mock_urlopen: MagicMock) -> None:
+ """Verify Authorization, Content-Type, and User-Agent headers are set."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials(access_token="my-token")
+
+ platform_request("/api/test", credentials=creds)
+
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.get_header("Authorization") == "Bearer my-token"
+ assert request_obj.get_header("Content-type") == "application/json"
+ assert "osmosis-cli/" in request_obj.get_header("User-agent")
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_custom_headers_are_merged(self, mock_urlopen: MagicMock) -> None:
+ """Verify additional headers are merged into the request."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials()
+
+ platform_request(
+ "/api/test",
+ headers={"X-Custom": "value123"},
+ credentials=creds,
+ )
+
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.get_header("X-custom") == "value123"
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_get_request_has_no_body(self, mock_urlopen: MagicMock) -> None:
+ """Verify GET requests send no request body."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials()
+
+ platform_request("/api/test", method="GET", credentials=creds)
+
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.data is None
+ assert request_obj.get_method() == "GET"
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_post_request_sends_json_body(self, mock_urlopen: MagicMock) -> None:
+ """Verify POST requests with data encode body as JSON."""
+ mock_urlopen.return_value = _make_http_response({"created": True})
+ creds = _make_credentials()
+ payload = {"host": "10.0.0.1", "port": 8080}
+
+ platform_request(
+ "/api/register", method="POST", data=payload, credentials=creds
+ )
+
+ request_obj = mock_urlopen.call_args[0][0]
+ assert request_obj.get_method() == "POST"
+ body = json.loads(request_obj.data.decode())
+ assert body == {"host": "10.0.0.1", "port": 8080}
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_timeout_is_passed_to_urlopen(self, mock_urlopen: MagicMock) -> None:
+ """Verify the timeout parameter is forwarded to urlopen."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials()
+
+ platform_request("/api/test", timeout=5.0, credentials=creds)
+
+ _, kwargs = mock_urlopen.call_args
+ assert kwargs["timeout"] == 5.0
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_default_timeout_is_30(self, mock_urlopen: MagicMock) -> None:
+ """Verify default timeout is 30 seconds."""
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+ creds = _make_credentials()
+
+ platform_request("/api/test", credentials=creds)
+
+ _, kwargs = mock_urlopen.call_args
+ assert kwargs["timeout"] == 30.0
+
+ # -------------------------------------------------------------------------
+ # Successful Response Handling
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_returns_parsed_json_response(self, mock_urlopen: MagicMock) -> None:
+ """Verify the response JSON is parsed and returned."""
+ expected = {"id": "srv_1", "status": "healthy"}
+ mock_urlopen.return_value = _make_http_response(expected)
+ creds = _make_credentials()
+
+ result = platform_request("/api/test", credentials=creds)
+
+ assert result == expected
+
+ # -------------------------------------------------------------------------
+ # HTTP Error Handling
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client._handle_401_and_cleanup")
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_401_triggers_cleanup_and_raises(
+ self, mock_urlopen: MagicMock, mock_cleanup: MagicMock
+ ) -> None:
+ """Verify 401 response triggers credential cleanup."""
+ mock_urlopen.side_effect = _make_http_error(401, "Unauthorized")
+ mock_cleanup.side_effect = AuthenticationExpiredError("expired")
+ creds = _make_credentials()
+
+ with pytest.raises(AuthenticationExpiredError):
+ platform_request("/api/test", credentials=creds)
+
+ mock_cleanup.assert_called_once()
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_non_401_http_error_raises_platform_api_error(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify non-401 HTTP errors raise PlatformAPIError with status code."""
+ mock_urlopen.side_effect = _make_http_error(500, "Internal Server Error")
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ assert exc_info.value.status_code == 500
+ assert "HTTP 500" in str(exc_info.value)
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_http_error_includes_response_body_in_message(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify the error response body is captured in the PlatformAPIError message."""
+ error_body = '{"detail": "Agent not found"}'
+ mock_urlopen.side_effect = _make_http_error(404, error_body)
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ assert "Agent not found" in str(exc_info.value)
+ assert exc_info.value.status_code == 404
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_http_error_truncates_long_response_body(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify response bodies longer than 500 chars are truncated."""
+ long_body = "x" * 600
+ mock_urlopen.side_effect = _make_http_error(500, long_body)
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ msg = str(exc_info.value)
+ assert "(truncated)" in msg
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_http_error_with_empty_body(self, mock_urlopen: MagicMock) -> None:
+ """Verify HTTP errors with empty body still produce a useful message."""
+ mock_urlopen.side_effect = _make_http_error(502, "")
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ assert "HTTP 502" in str(exc_info.value)
+ assert exc_info.value.status_code == 502
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_http_error_with_unreadable_body_does_not_crash(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify HTTP error handling is robust when body read fails."""
+ err = _make_http_error(503, "error")
+ # Make the read() call raise, simulating a broken stream
+ err.read = MagicMock(side_effect=OSError("stream broken"))
+ mock_urlopen.side_effect = err
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ # Should still get a PlatformAPIError with the status code
+ assert exc_info.value.status_code == 503
+ assert "HTTP 503" in str(exc_info.value)
+
+ # -------------------------------------------------------------------------
+ # URLError (Connection Errors)
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_url_error_raises_platform_api_error(self, mock_urlopen: MagicMock) -> None:
+ """Verify URLError is wrapped in PlatformAPIError."""
+ mock_urlopen.side_effect = URLError("Name resolution failed")
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError, match="Connection error"):
+ platform_request("/api/test", credentials=creds)
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_url_error_does_not_include_status_code(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify URLError-based PlatformAPIError has no status code."""
+ mock_urlopen.side_effect = URLError("Connection refused")
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError) as exc_info:
+ platform_request("/api/test", credentials=creds)
+
+ assert exc_info.value.status_code is None
+
+ # -------------------------------------------------------------------------
+ # JSON Decode Errors
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ def test_invalid_json_response_raises_platform_api_error(
+ self, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify non-JSON responses raise PlatformAPIError."""
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = b"Not JSON"
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+ mock_urlopen.return_value = mock_resp
+ creds = _make_credentials()
+
+ with pytest.raises(PlatformAPIError, match="Invalid JSON response"):
+ platform_request("/api/test", credentials=creds)
+
+ # -------------------------------------------------------------------------
+ # Integration-style: Credential flow
+ # -------------------------------------------------------------------------
+
+ @patch("osmosis_ai.auth.platform_client.urlopen")
+ @patch("osmosis_ai.auth.platform_client.load_credentials")
+ def test_explicit_none_credentials_triggers_load(
+ self, mock_load: MagicMock, mock_urlopen: MagicMock
+ ) -> None:
+ """Verify passing credentials=None explicitly still loads from storage."""
+ creds = _make_credentials()
+ mock_load.return_value = creds
+ mock_urlopen.return_value = _make_http_response({"ok": True})
+
+ platform_request("/api/test", credentials=None)
+
+ mock_load.assert_called_once()
+
+ @patch("osmosis_ai.auth.platform_client.load_credentials")
+ def test_load_returns_none_raises_auth_error(self, mock_load: MagicMock) -> None:
+ """Verify AuthenticationExpiredError when load_credentials returns None."""
+ mock_load.return_value = None
+
+ with pytest.raises(AuthenticationExpiredError, match="No valid credentials"):
+ platform_request("/api/test", credentials=None)
diff --git a/tests/test_registration.py b/tests/test_registration.py
new file mode 100644
index 00000000..d37763bd
--- /dev/null
+++ b/tests/test_registration.py
@@ -0,0 +1,513 @@
+"""Tests for osmosis_ai.rollout.server.registration."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from osmosis_ai.rollout.network import PublicIPDetectionError
+from osmosis_ai.rollout.server.registration import (
+ RegistrationResult,
+ get_public_ip,
+ get_report_host,
+ print_registration_result,
+ register_with_platform,
+)
+
+# =============================================================================
+# Helper: Create mock WorkspaceCredentials
+# =============================================================================
+
+
+def _make_mock_credentials() -> MagicMock:
+ """Create a mock WorkspaceCredentials with a fake access_token."""
+ creds = MagicMock()
+ creds.access_token = "fake-token-abc123"
+ return creds
+
+
+# =============================================================================
+# RegistrationResult Tests
+# =============================================================================
+
+
+class TestRegistrationResult:
+ """Tests for the RegistrationResult dataclass."""
+
+ def test_default_values(self) -> None:
+ """Verify default field values on a minimal RegistrationResult."""
+ result = RegistrationResult(success=True)
+ assert result.success is True
+ assert result.server_id is None
+ assert result.status == "unknown"
+ assert result.error is None
+ assert result.server_info is None
+
+ def test_is_healthy_when_status_is_healthy(self) -> None:
+ """Verify is_healthy returns True only when status is 'healthy'."""
+ result = RegistrationResult(success=True, status="healthy")
+ assert result.is_healthy is True
+
+ def test_is_healthy_returns_false_for_other_statuses(self) -> None:
+ """Verify is_healthy returns False for non-healthy statuses."""
+ for status in ("unknown", "error", "unhealthy", "pending", ""):
+ result = RegistrationResult(success=True, status=status)
+ assert result.is_healthy is False, f"Expected False for status={status!r}"
+
+ def test_full_construction(self) -> None:
+ """Verify all fields can be set at construction time."""
+ info = {"active_rollouts": 3, "agent_loop": "my_agent"}
+ result = RegistrationResult(
+ success=True,
+ server_id="srv_123",
+ status="healthy",
+ error=None,
+ server_info=info,
+ )
+ assert result.server_id == "srv_123"
+ assert result.server_info == info
+
+
+# =============================================================================
+# get_public_ip Tests
+# =============================================================================
+
+
+class TestGetPublicIP:
+ """Tests for the get_public_ip function."""
+
+ @patch("osmosis_ai.rollout.server.registration.detect_public_ip")
+ def test_returns_detected_ip(self, mock_detect: MagicMock) -> None:
+ """Verify get_public_ip delegates to detect_public_ip and returns its result."""
+ mock_detect.return_value = "54.200.100.50"
+ assert get_public_ip() == "54.200.100.50"
+ mock_detect.assert_called_once()
+
+ @patch("osmosis_ai.rollout.server.registration.detect_public_ip")
+ def test_propagates_detection_error(self, mock_detect: MagicMock) -> None:
+ """Verify PublicIPDetectionError propagates from get_public_ip."""
+ mock_detect.side_effect = PublicIPDetectionError("all methods failed")
+ with pytest.raises(PublicIPDetectionError, match="all methods failed"):
+ get_public_ip()
+
+
+# =============================================================================
+# get_report_host Tests
+# =============================================================================
+
+
+class TestGetReportHost:
+ """Tests for the get_report_host function."""
+
+ @patch("osmosis_ai.rollout.server.registration.get_public_ip")
+ def test_returns_public_ip_for_all_interfaces(self, mock_ip: MagicMock) -> None:
+ """Verify 0.0.0.0 triggers public IP detection."""
+ mock_ip.return_value = "203.0.113.10"
+ assert get_report_host("0.0.0.0") == "203.0.113.10"
+ mock_ip.assert_called_once()
+
+ def test_returns_host_as_is_for_specific_address(self) -> None:
+ """Verify a specific IP/hostname is returned unchanged."""
+ assert get_report_host("10.0.0.5") == "10.0.0.5"
+ assert get_report_host("my-server.example.com") == "my-server.example.com"
+ assert get_report_host("127.0.0.1") == "127.0.0.1"
+
+ @patch("osmosis_ai.rollout.server.registration.get_public_ip")
+ def test_propagates_error_for_all_interfaces(self, mock_ip: MagicMock) -> None:
+ """Verify PublicIPDetectionError propagates when host is 0.0.0.0."""
+ mock_ip.side_effect = PublicIPDetectionError("detection failed")
+ with pytest.raises(PublicIPDetectionError, match="detection failed"):
+ get_report_host("0.0.0.0")
+
+
+# =============================================================================
+# register_with_platform Tests
+# =============================================================================
+
+
+class TestRegisterWithPlatform:
+ """Tests for the register_with_platform function."""
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_healthy_registration(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify a successful healthy registration returns expected result."""
+ mock_host.return_value = "54.200.100.1"
+ mock_request.return_value = {
+ "id": "srv_abc",
+ "status": "healthy",
+ "health_check_result": {
+ "server_info": {"active_rollouts": 0, "agent_loop": "my_agent"},
+ },
+ }
+
+ result = register_with_platform(
+ host="0.0.0.0",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is True
+ assert result.is_healthy is True
+ assert result.server_id == "srv_abc"
+ assert result.server_info == {"active_rollouts": 0, "agent_loop": "my_agent"}
+ assert result.error is None
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_registration_succeeded_but_health_check_failed(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify registration success with health check failure returns appropriate result."""
+ mock_host.return_value = "54.200.100.1"
+ mock_request.return_value = {
+ "id": "srv_xyz",
+ "status": "unhealthy",
+ "health_check_result": {
+ "error": "Connection refused",
+ },
+ }
+
+ result = register_with_platform(
+ host="0.0.0.0",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is True
+ assert result.is_healthy is False
+ assert result.server_id == "srv_xyz"
+ assert result.status == "unhealthy"
+ assert result.error == "Connection refused"
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_registration_sends_correct_data_without_api_key(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify the correct registration payload is sent when no api_key is given."""
+ mock_host.return_value = "10.0.0.5"
+ mock_request.return_value = {
+ "id": "srv_1",
+ "status": "healthy",
+ "health_check_result": {},
+ }
+
+ creds = _make_mock_credentials()
+ register_with_platform(
+ host="10.0.0.5",
+ port=9090,
+ agent_loop_name="test_agent",
+ credentials=creds,
+ )
+
+ mock_request.assert_called_once_with(
+ "/api/cli/rollout-server/register",
+ method="POST",
+ data={
+ "host": "10.0.0.5",
+ "port": 9090,
+ "agent_loop_name": "test_agent",
+ },
+ timeout=15.0,
+ credentials=creds,
+ )
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_registration_sends_api_key_when_provided(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify api_key is included in the registration payload when provided."""
+ mock_host.return_value = "10.0.0.5"
+ mock_request.return_value = {
+ "id": "srv_1",
+ "status": "healthy",
+ "health_check_result": {},
+ }
+
+ creds = _make_mock_credentials()
+ register_with_platform(
+ host="10.0.0.5",
+ port=9090,
+ agent_loop_name="test_agent",
+ credentials=creds,
+ api_key="secret-key-456",
+ )
+
+ call_data = mock_request.call_args[1]["data"]
+ assert call_data["api_key"] == "secret-key-456"
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_ip_detection_failure_returns_error_result(
+ self, mock_host: MagicMock
+ ) -> None:
+ """Verify IP detection failure returns a failed RegistrationResult."""
+ mock_host.side_effect = PublicIPDetectionError("no IP")
+
+ result = register_with_platform(
+ host="0.0.0.0",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is False
+ assert result.status == "error"
+ assert "public IP" in result.error
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_authentication_expired_returns_error_result(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify AuthenticationExpiredError is caught and returned as error result."""
+ from osmosis_ai.auth.platform_client import AuthenticationExpiredError
+
+ mock_host.return_value = "10.0.0.5"
+ mock_request.side_effect = AuthenticationExpiredError("session expired")
+
+ result = register_with_platform(
+ host="10.0.0.5",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is False
+ assert result.status == "error"
+ assert "session expired" in result.error
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_platform_api_error_returns_error_result(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify PlatformAPIError is caught and returned as error result."""
+ from osmosis_ai.auth.platform_client import PlatformAPIError
+
+ mock_host.return_value = "10.0.0.5"
+ mock_request.side_effect = PlatformAPIError("HTTP 500", status_code=500)
+
+ result = register_with_platform(
+ host="10.0.0.5",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is False
+ assert result.status == "error"
+ assert "HTTP 500" in result.error
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_unexpected_exception_returns_error_result(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify unexpected exceptions are caught and returned as error result."""
+ mock_host.return_value = "10.0.0.5"
+ mock_request.side_effect = RuntimeError("unexpected network issue")
+
+ result = register_with_platform(
+ host="10.0.0.5",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is False
+ assert result.status == "error"
+ assert "unexpected network issue" in result.error
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_missing_status_defaults_to_unknown(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify missing status in platform response defaults to 'unknown'."""
+ mock_host.return_value = "10.0.0.5"
+ mock_request.return_value = {
+ "id": "srv_1",
+ # No "status" key
+ "health_check_result": {},
+ }
+
+ result = register_with_platform(
+ host="10.0.0.5",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ # Registration succeeded, but status is unknown -> non-healthy branch
+ assert result.success is True
+ assert result.status == "unknown"
+ assert result.is_healthy is False
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ @patch("osmosis_ai.auth.platform_client.platform_request")
+ def test_health_check_failed_default_error_message(
+ self, mock_request: MagicMock, mock_host: MagicMock
+ ) -> None:
+ """Verify default error message when health_check_result has no 'error' key."""
+ mock_host.return_value = "10.0.0.5"
+ mock_request.return_value = {
+ "id": "srv_1",
+ "status": "unhealthy",
+ "health_check_result": {}, # No "error" key
+ }
+
+ result = register_with_platform(
+ host="10.0.0.5",
+ port=8080,
+ agent_loop_name="my_agent",
+ credentials=_make_mock_credentials(),
+ )
+
+ assert result.success is True
+ assert result.error == "Health check failed"
+
+
+# =============================================================================
+# print_registration_result Tests
+# =============================================================================
+
+
+class TestPrintRegistrationResult:
+ """Tests for the print_registration_result function."""
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_healthy_output(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify output for a healthy registration result."""
+ mock_host.return_value = "54.200.100.1"
+ result = RegistrationResult(
+ success=True,
+ server_id="srv_1",
+ status="healthy",
+ server_info={"active_rollouts": 2},
+ )
+
+ print_registration_result(
+ result, host="0.0.0.0", port=8080, agent_loop_name="my_agent"
+ )
+
+ output = capsys.readouterr().out
+ assert "[OK] Registered with Osmosis Platform" in output
+ assert "my_agent" in output
+ assert "54.200.100.1:8080" in output
+ assert "Active rollouts: 2" in output
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_healthy_output_without_server_info(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify output for healthy result when server_info is None."""
+ mock_host.return_value = "54.200.100.1"
+ result = RegistrationResult(
+ success=True,
+ server_id="srv_1",
+ status="healthy",
+ server_info=None,
+ )
+
+ print_registration_result(
+ result, host="0.0.0.0", port=8080, agent_loop_name="my_agent"
+ )
+
+ output = capsys.readouterr().out
+ assert "[OK] Registered with Osmosis Platform" in output
+ assert "Active rollouts" not in output
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_unhealthy_output(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify output for registration success but health check failure."""
+ mock_host.return_value = "10.0.0.5"
+ result = RegistrationResult(
+ success=True,
+ server_id="srv_2",
+ status="unhealthy",
+ error="Connection refused",
+ )
+
+ print_registration_result(
+ result, host="10.0.0.5", port=9090, agent_loop_name="test_agent"
+ )
+
+ output = capsys.readouterr().out
+ assert "[WARNING] Registered but health check failed" in output
+ assert "test_agent" in output
+ assert "10.0.0.5:9090" in output
+ assert "Connection refused" in output
+ assert "The server will continue running." in output
+ assert "Tip: Use a VM with public IP" in output
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_failed_registration_output(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify output for a completely failed registration."""
+ mock_host.return_value = "10.0.0.5"
+ result = RegistrationResult(
+ success=False,
+ status="error",
+ error="Authentication expired",
+ )
+
+ print_registration_result(
+ result, host="10.0.0.5", port=8080, agent_loop_name="agent"
+ )
+
+ output = capsys.readouterr().out
+ assert "[WARNING] Failed to register with Platform" in output
+ assert "Authentication expired" in output
+ assert "The server will continue running without registration." in output
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_ip_detection_failure_falls_back_to_original_host(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify print_registration_result falls back to original host on IP detection error."""
+ mock_host.side_effect = PublicIPDetectionError("cannot detect IP")
+ result = RegistrationResult(
+ success=True,
+ server_id="srv_1",
+ status="healthy",
+ )
+
+ print_registration_result(
+ result, host="0.0.0.0", port=8080, agent_loop_name="agent"
+ )
+
+ output = capsys.readouterr().out
+ # Should fall back to original host "0.0.0.0" instead of crashing
+ assert "0.0.0.0:8080" in output
+ assert "[OK] Registered with Osmosis Platform" in output
+
+ @patch("osmosis_ai.rollout.server.registration.get_report_host")
+ def test_healthy_output_with_zero_active_rollouts(
+ self, mock_host: MagicMock, capsys: pytest.CaptureFixture[str]
+ ) -> None:
+ """Verify active rollouts line shows 0 correctly."""
+ mock_host.return_value = "10.0.0.5"
+ result = RegistrationResult(
+ success=True,
+ status="healthy",
+ server_info={"active_rollouts": 0},
+ )
+
+ print_registration_result(
+ result, host="10.0.0.5", port=8080, agent_loop_name="agent"
+ )
+
+ output = capsys.readouterr().out
+ assert "Active rollouts: 0" in output
diff --git a/tests/test_rollout_server_app.py b/tests/test_rollout_server_app.py
new file mode 100644
index 00000000..26180e59
--- /dev/null
+++ b/tests/test_rollout_server_app.py
@@ -0,0 +1,1037 @@
+"""Tests for osmosis_ai.rollout.server.app — focusing on uncovered code paths.
+
+This module tests:
+- _extract_bearer_token() parsing logic
+- lifespan() startup/shutdown lifecycle (cleanup task, callbacks, registration)
+- Concurrency control via semaphore (max_concurrent)
+- Idempotency with idempotency_key vs rollout_id
+- Error propagation when get_tools() raises
+- Platform registration background task behaviour
+- on_startup / on_shutdown lifecycle callbacks
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from osmosis_ai.rollout import (
+ OpenAIFunctionToolSchema,
+ RolloutAgentLoop,
+ RolloutContext,
+ RolloutMetrics,
+ RolloutRequest,
+ RolloutResult,
+ create_app,
+)
+from osmosis_ai.rollout.server.app import _extract_bearer_token
+
+# Import FastAPI / httpx test utilities
+try:
+ from fastapi.testclient import TestClient
+ from httpx import ASGITransport, AsyncClient
+
+ FASTAPI_AVAILABLE = True
+except ImportError:
+ FASTAPI_AVAILABLE = False
+
+pytestmark = pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
+
+
+# =============================================================================
+# Helpers — reusable agent loop stubs
+# =============================================================================
+
+
+class SimpleAgentLoop(RolloutAgentLoop):
+ """Agent loop that completes immediately."""
+
+ name = "simple_agent"
+
+ def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
+ self._tools = tools or []
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return self._tools
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ return ctx.complete(list(ctx.request.messages))
+
+
+class SlowAgentLoop(RolloutAgentLoop):
+ """Agent loop that sleeps for a configurable duration before completing."""
+
+ name = "slow_agent"
+
+ def __init__(self, delay: float = 0.5):
+ self._delay = delay
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ await asyncio.sleep(self._delay)
+ return ctx.complete(list(ctx.request.messages))
+
+
+class FailingGetToolsAgentLoop(RolloutAgentLoop):
+ """Agent loop whose get_tools() raises an exception."""
+
+ name = "failing_get_tools_agent"
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ raise ValueError("Tools unavailable")
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ return ctx.complete([])
+
+
+def _make_rollout_payload(
+ rollout_id: str = "test-123",
+ idempotency_key: str | None = None,
+ **overrides: Any,
+) -> dict[str, Any]:
+ """Build a minimal valid RolloutRequest JSON payload."""
+ payload: dict[str, Any] = {
+ "rollout_id": rollout_id,
+ "server_url": "http://localhost:8080",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "completion_params": {"temperature": 0.7},
+ }
+ if idempotency_key is not None:
+ payload["idempotency_key"] = idempotency_key
+ payload.update(overrides)
+ return payload
+
+
+def _mock_llm_client() -> MagicMock:
+ """Create a mocked OsmosisLLMClient usable as an async context manager."""
+ mock = AsyncMock()
+ mock.__aenter__ = AsyncMock(return_value=mock)
+ mock.__aexit__ = AsyncMock(return_value=None)
+ mock.complete_rollout = AsyncMock()
+ mock.get_metrics = MagicMock(return_value=RolloutMetrics())
+ return mock
+
+
+# =============================================================================
+# _extract_bearer_token() tests
+# =============================================================================
+
+
+class TestExtractBearerToken:
+ """Tests for the _extract_bearer_token helper function."""
+
+ def test_returns_none_for_none_input(self) -> None:
+ """None header should yield None."""
+ assert _extract_bearer_token(None) is None # type: ignore[arg-type]
+
+ def test_returns_none_for_empty_string(self) -> None:
+ """Empty string should yield None."""
+ assert _extract_bearer_token("") is None
+
+ def test_returns_none_for_whitespace_only(self) -> None:
+ """Whitespace-only header should yield None."""
+ assert _extract_bearer_token(" ") is None
+
+ def test_extracts_bearer_prefix(self) -> None:
+ """Standard 'Bearer ' format should return ."""
+ assert _extract_bearer_token("Bearer my-secret-token") == "my-secret-token"
+
+ def test_bearer_prefix_is_case_insensitive(self) -> None:
+ """The 'bearer' prefix should be matched case-insensitively."""
+ assert _extract_bearer_token("bearer abc123") == "abc123"
+ assert _extract_bearer_token("BEARER xyz789") == "xyz789"
+ assert _extract_bearer_token("BeArEr mixed") == "mixed"
+
+ def test_raw_token_fallback(self) -> None:
+ """A single token without Bearer prefix should be returned as-is."""
+ assert _extract_bearer_token("plain-token") == "plain-token"
+
+ def test_strips_surrounding_whitespace(self) -> None:
+ """Leading/trailing whitespace around the header should be stripped."""
+ assert _extract_bearer_token(" Bearer tok ") == "tok"
+
+ def test_more_than_two_parts_returns_raw(self) -> None:
+ """If there are more than two parts, fall back to returning the full string."""
+ result = _extract_bearer_token("Bearer token extra")
+ # Three parts -- does not match the 'len(parts) == 2' branch
+ assert result == "Bearer token extra"
+
+ def test_non_bearer_prefix_with_two_parts_returns_raw(self) -> None:
+ """Two-part header with a non-bearer prefix returns raw string."""
+ result = _extract_bearer_token("Basic dXNlcjpwYXNz")
+ assert result == "Basic dXNlcjpwYXNz"
+
+
+# =============================================================================
+# Lifespan -- on_startup / on_shutdown callbacks
+# =============================================================================
+
+
+class TestLifespanCallbacks:
+ """Tests that on_startup and on_shutdown callbacks are invoked.
+
+ Uses TestClient as a context manager which properly triggers lifespan events.
+ """
+
+ def test_on_startup_called_during_lifespan(self) -> None:
+ """on_startup callback should be invoked when the app starts."""
+ started = False
+
+ async def my_startup() -> None:
+ nonlocal started
+ started = True
+
+ app = create_app(SimpleAgentLoop(), on_startup=my_startup)
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ assert started, "on_startup was not called"
+
+ def test_on_shutdown_called_during_lifespan(self) -> None:
+ """on_shutdown callback should be invoked when the app shuts down."""
+ stopped = False
+
+ async def my_shutdown() -> None:
+ nonlocal stopped
+ stopped = True
+
+ app = create_app(SimpleAgentLoop(), on_shutdown=my_shutdown)
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ # After the context manager exits, shutdown should have fired
+ assert stopped, "on_shutdown was not called"
+
+ def test_both_callbacks_called_in_order(self) -> None:
+ """on_startup fires before requests; on_shutdown fires after exit."""
+ events: list[str] = []
+
+ async def startup() -> None:
+ events.append("start")
+
+ async def shutdown() -> None:
+ events.append("stop")
+
+ app = create_app(SimpleAgentLoop(), on_startup=startup, on_shutdown=shutdown)
+ with TestClient(app) as client:
+ client.get("/health")
+ events.append("request")
+
+ assert events == ["start", "request", "stop"]
+
+
+# =============================================================================
+# Lifespan -- cleanup task management
+# =============================================================================
+
+
+class TestLifespanCleanupTask:
+ """Tests that the cleanup task starts and stops with the lifespan."""
+
+ def test_cleanup_task_started_on_startup(self) -> None:
+ """The AppState cleanup task should be started during lifespan startup."""
+ app = create_app(SimpleAgentLoop())
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ # If we get here without errors, the cleanup task started and was
+ # stopped cleanly during shutdown.
+
+
+# =============================================================================
+# Lifespan -- platform registration
+# =============================================================================
+
+
+class TestLifespanRegistration:
+ """Tests for the platform registration background task in lifespan."""
+
+ def test_registration_task_not_started_without_credentials(self) -> None:
+ """Registration should be skipped when credentials are None."""
+ app = create_app(SimpleAgentLoop())
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+
+ def test_registration_not_started_when_host_is_none(self) -> None:
+ """Registration should be skipped if server_host is None."""
+ mock_creds = MagicMock()
+ app = create_app(
+ SimpleAgentLoop(),
+ credentials=mock_creds,
+ server_host=None,
+ server_port=9999,
+ )
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+
+ def test_registration_not_started_when_port_is_none(self) -> None:
+ """Registration should be skipped if server_port is None."""
+ mock_creds = MagicMock()
+ app = create_app(
+ SimpleAgentLoop(),
+ credentials=mock_creds,
+ server_host="0.0.0.0",
+ server_port=None,
+ )
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+
+ async def test_registration_task_started_with_credentials(self) -> None:
+ """Registration background task should be created when credentials are given."""
+ mock_creds = MagicMock()
+ mock_register = MagicMock(return_value={"status": "ok"})
+ mock_print = MagicMock()
+
+ with (
+ patch(
+ "osmosis_ai.rollout.server.registration.register_with_platform",
+ mock_register,
+ ),
+ patch(
+ "osmosis_ai.rollout.server.registration.print_registration_result",
+ mock_print,
+ ),
+ ):
+ app = create_app(
+ SimpleAgentLoop(),
+ credentials=mock_creds,
+ server_host="0.0.0.0",
+ server_port=9999,
+ )
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ # Give the registration background task time to run
+ await asyncio.sleep(0.5)
+
+ async def test_registration_error_handled_gracefully(self) -> None:
+ """If the registration raises an error, shutdown should not crash."""
+ mock_creds = MagicMock()
+
+ with (
+ patch(
+ "osmosis_ai.rollout.server.registration.register_with_platform",
+ side_effect=Exception("Registration failed"),
+ ),
+ patch(
+ "osmosis_ai.rollout.server.registration.print_registration_result",
+ MagicMock(),
+ ),
+ ):
+ app = create_app(
+ SimpleAgentLoop(),
+ credentials=mock_creds,
+ server_host="0.0.0.0",
+ server_port=9998,
+ )
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ # Allow time for registration attempt
+ await asyncio.sleep(0.5)
+ # TestClient should exit cleanly even though registration errored
+
+
+# =============================================================================
+# Concurrency control (max_concurrent)
+# =============================================================================
+
+
+class TestConcurrencyControl:
+ """Tests for the semaphore-based max_concurrent limit."""
+
+ async def test_semaphore_limits_concurrent_rollouts(self) -> None:
+ """Only max_concurrent rollouts should run simultaneously."""
+ max_concurrent = 2
+ agent = SlowAgentLoop(delay=0.3)
+ app = create_app(agent, max_concurrent=max_concurrent)
+
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ # Submit 4 rollouts
+ for i in range(4):
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id=f"conc-{i}"),
+ )
+ assert resp.status_code == 202
+
+ # Allow some time for background tasks
+ await asyncio.sleep(0.1)
+
+ # Check health to see active rollouts
+ health = await client.get("/health")
+ data = health.json()
+ # There should be some active rollouts queued up
+ assert data["active_rollouts"] >= 0
+
+ # Wait for all to complete
+ await asyncio.sleep(1.5)
+
+ async def test_max_concurrent_one_serializes_rollouts(self) -> None:
+ """With max_concurrent=1 rollouts must execute one at a time."""
+ execution_order: list[str] = []
+
+ class OrderTrackingAgent(RolloutAgentLoop):
+ name = "order_agent"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ rid = ctx.request.rollout_id
+ execution_order.append(f"start-{rid}")
+ await asyncio.sleep(0.05)
+ execution_order.append(f"end-{rid}")
+ return ctx.complete(list(ctx.request.messages))
+
+ app = create_app(OrderTrackingAgent(), max_concurrent=1)
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="r1"),
+ )
+ await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="r2"),
+ )
+ # Wait for both to finish
+ await asyncio.sleep(0.5)
+
+ # With max_concurrent=1, r1 should fully complete before r2 starts
+ assert "start-r1" in execution_order
+ assert "end-r1" in execution_order
+ assert "start-r2" in execution_order
+ assert "end-r2" in execution_order
+ r1_end_idx = execution_order.index("end-r1")
+ r2_start_idx = execution_order.index("start-r2")
+ assert r1_end_idx < r2_start_idx, (
+ f"Expected r1 to finish before r2 starts, but got: {execution_order}"
+ )
+
+
+# =============================================================================
+# Idempotency with idempotency_key
+# =============================================================================
+
+
+class TestIdempotencyKey:
+ """Tests for idempotency_key handling in /v1/rollout/init."""
+
+ async def test_idempotency_key_used_when_provided(self) -> None:
+ """Two requests with different rollout_ids but same idempotency_key
+ should be treated as duplicates."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp1 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="id-1", idempotency_key="shared-key"
+ ),
+ )
+ resp2 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="id-2", idempotency_key="shared-key"
+ ),
+ )
+
+ assert resp1.status_code == 202
+ assert resp2.status_code == 202
+ # Both should return the same rollout_id (from the leader request)
+ assert resp1.json()["rollout_id"] == resp2.json()["rollout_id"]
+
+ async def test_different_idempotency_keys_are_independent(self) -> None:
+ """Requests with different idempotency_keys should be treated separately."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp1 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="id-a", idempotency_key="key-a"
+ ),
+ )
+ resp2 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="id-b", idempotency_key="key-b"
+ ),
+ )
+
+ assert resp1.status_code == 202
+ assert resp2.status_code == 202
+ assert resp1.json()["rollout_id"] != resp2.json()["rollout_id"]
+
+ async def test_fallback_to_rollout_id_when_no_idempotency_key(self) -> None:
+ """Without idempotency_key, rollout_id is the idempotency key."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp1 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="dup-id"),
+ )
+ resp2 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="dup-id"),
+ )
+
+ assert resp1.status_code == 202
+ assert resp2.status_code == 202
+ assert resp1.json()["rollout_id"] == resp2.json()["rollout_id"]
+
+
+# =============================================================================
+# Error handling in init endpoint
+# =============================================================================
+
+
+class TestInitEndpointErrorHandling:
+ """Tests for error paths in the /v1/rollout/init handler."""
+
+ async def test_get_tools_error_returns_500(self) -> None:
+ """If get_tools() raises, the endpoint should return 500."""
+ app = create_app(FailingGetToolsAgentLoop())
+ # raise_app_exceptions=False so httpx returns the HTTP 500 response
+ # instead of propagating the exception to the test.
+ transport = ASGITransport(app=app, raise_app_exceptions=False)
+ async with AsyncClient(transport=transport, base_url="http://test") as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="fail-tools"),
+ )
+ assert resp.status_code == 500
+
+ async def test_get_tools_error_clears_init_record(self) -> None:
+ """After get_tools() fails, a retry with the same ID should succeed
+ (the init record should have been cleaned up)."""
+
+ call_count = 0
+
+ class FailOnceThenSucceedAgent(RolloutAgentLoop):
+ name = "fail_once_agent"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ nonlocal call_count
+ call_count += 1
+ if call_count == 1:
+ raise RuntimeError("Transient error")
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ return ctx.complete([])
+
+ app = create_app(FailOnceThenSucceedAgent())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app, raise_app_exceptions=False)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ # First request fails
+ resp1 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="retry-me"),
+ )
+ assert resp1.status_code == 500
+
+ # Retry should succeed (init record was cleared)
+ resp2 = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="retry-me"),
+ )
+ assert resp2.status_code == 202
+
+ async def test_infrastructure_error_marks_completed(self) -> None:
+ """When OsmosisLLMClient raises during setup, the key should still be
+ marked completed so resources are cleaned up."""
+ app = create_app(SimpleAgentLoop())
+
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ mock = AsyncMock()
+ mock.__aenter__ = AsyncMock(side_effect=ConnectionError("Cannot connect"))
+ mock.__aexit__ = AsyncMock(return_value=None)
+ MockCls.return_value = mock
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="infra-err"),
+ )
+ assert resp.status_code == 202
+
+ # Give background task time to fail and clean up
+ await asyncio.sleep(0.2)
+
+ health = await client.get("/health")
+ data = health.json()
+ # Should have moved to completed (not stuck in active)
+ assert data["active_rollouts"] == 0
+
+
+# =============================================================================
+# Background rollout task
+# =============================================================================
+
+
+class TestBackgroundRolloutTask:
+ """Tests for the run_rollout() background task created in init_rollout."""
+
+ async def test_successful_rollout_calls_complete_rollout(self) -> None:
+ """A successful agent run should call llm.complete_rollout with COMPLETED."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="bg-ok"),
+ )
+ assert resp.status_code == 202
+ await asyncio.sleep(0.2)
+
+ mock_client.complete_rollout.assert_called_once()
+ kwargs = mock_client.complete_rollout.call_args[1]
+ assert kwargs["status"] == "COMPLETED"
+ assert kwargs["finish_reason"] == "stop"
+
+ async def test_agent_error_reports_error_status(self) -> None:
+ """When the agent raises, complete_rollout should be called with ERROR."""
+
+ class ErrorAgent(RolloutAgentLoop):
+ name = "error_agent"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ raise RuntimeError("Agent crashed")
+
+ app = create_app(ErrorAgent())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="bg-err"),
+ )
+ assert resp.status_code == 202
+ await asyncio.sleep(0.2)
+
+ mock_client.complete_rollout.assert_called_once()
+ kwargs = mock_client.complete_rollout.call_args[1]
+ assert kwargs["status"] == "ERROR"
+ assert "Agent crashed" in kwargs["error_message"]
+
+ async def test_rollout_uses_request_api_key_for_callback(self) -> None:
+ """The OsmosisLLMClient should be constructed with the request's api_key."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="api-key-cb", api_key="callback-secret"
+ ),
+ )
+ assert resp.status_code == 202
+ await asyncio.sleep(0.2)
+
+ # OsmosisLLMClient should have been called with the request's api_key
+ MockCls.assert_called_once()
+ assert MockCls.call_args.kwargs["api_key"] == "callback-secret"
+
+ async def test_rollout_passes_server_url_to_client(self) -> None:
+ """The OsmosisLLMClient should receive the server_url from the request."""
+ app = create_app(SimpleAgentLoop())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ resp = await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(
+ rollout_id="url-test",
+ server_url="http://traingate:9000",
+ ),
+ )
+ assert resp.status_code == 202
+ await asyncio.sleep(0.2)
+
+ MockCls.assert_called_once()
+ assert MockCls.call_args.kwargs["server_url"] == "http://traingate:9000"
+ assert MockCls.call_args.kwargs["rollout_id"] == "url-test"
+
+
+# =============================================================================
+# /platform/health endpoint edge cases
+# =============================================================================
+
+
+class TestPlatformHealthEdgeCases:
+ """Additional edge-case tests for /platform/health."""
+
+ def test_raw_token_without_bearer_prefix_accepted(self) -> None:
+ """A raw token (no 'Bearer' prefix) should still authenticate."""
+ api_key = "my-secret-key"
+ app = create_app(SimpleAgentLoop(), api_key=api_key)
+
+ with TestClient(app) as client:
+ resp = client.get(
+ "/platform/health",
+ headers={"Authorization": api_key},
+ )
+ assert resp.status_code == 200
+
+ def test_empty_authorization_header_returns_401(self) -> None:
+ """An empty Authorization header should be rejected."""
+ app = create_app(SimpleAgentLoop(), api_key="some-key")
+
+ with TestClient(app) as client:
+ resp = client.get(
+ "/platform/health",
+ headers={"Authorization": ""},
+ )
+ assert resp.status_code == 401
+
+ def test_platform_health_returns_agent_name_and_counts(self) -> None:
+ """Authenticated /platform/health should return agent info and counts."""
+ api_key = "check-key"
+ app = create_app(SimpleAgentLoop(), api_key=api_key)
+
+ with TestClient(app) as client:
+ resp = client.get(
+ "/platform/health",
+ headers={"Authorization": f"Bearer {api_key}"},
+ )
+ assert resp.status_code == 200
+ data = resp.json()
+ assert data["status"] == "healthy"
+ assert data["agent_loop"] == "simple_agent"
+ assert "active_rollouts" in data
+ assert "completed_rollouts" in data
+
+
+# =============================================================================
+# /v1/rollout/init API key auth edge cases
+# =============================================================================
+
+
+class TestInitApiKeyEdgeCases:
+ """Edge-case tests for API key authentication on /v1/rollout/init."""
+
+ def test_raw_token_accepted_on_init(self) -> None:
+ """A raw API key without 'Bearer' prefix should authenticate init."""
+ api_key = "my-secret-key"
+ app = create_app(SimpleAgentLoop(), api_key=api_key)
+
+ with TestClient(app) as client:
+ resp = client.post(
+ "/v1/rollout/init",
+ headers={"Authorization": api_key},
+ json=_make_rollout_payload(rollout_id="raw-auth"),
+ )
+ assert resp.status_code == 202
+
+ def test_empty_authorization_header_rejected_on_init(self) -> None:
+ """An empty Authorization header should be rejected on init."""
+ app = create_app(SimpleAgentLoop(), api_key="some-key")
+
+ with TestClient(app) as client:
+ resp = client.post(
+ "/v1/rollout/init",
+ headers={"Authorization": ""},
+ json=_make_rollout_payload(rollout_id="empty-auth"),
+ )
+ assert resp.status_code == 401
+
+ def test_no_auth_header_at_all_rejected(self) -> None:
+ """Missing Authorization header should be rejected when api_key is set."""
+ app = create_app(SimpleAgentLoop(), api_key="some-key")
+
+ with TestClient(app) as client:
+ resp = client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="no-auth"),
+ )
+ assert resp.status_code == 401
+
+ def test_no_api_key_configured_allows_unauthenticated(self) -> None:
+ """When api_key is not set, requests without auth should succeed."""
+ app = create_app(SimpleAgentLoop())
+
+ with TestClient(app) as client:
+ resp = client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="no-auth-needed"),
+ )
+ assert resp.status_code == 202
+
+
+# =============================================================================
+# /health endpoint details
+# =============================================================================
+
+
+class TestHealthEndpoint:
+ """Tests for the /health endpoint response shape."""
+
+ async def test_health_returns_correct_counts(self) -> None:
+ """Health endpoint should report accurate active and completed counts."""
+ agent = SlowAgentLoop(delay=0.3)
+ app = create_app(agent)
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ # Initially zero
+ health = await client.get("/health")
+ data = health.json()
+ assert data["active_rollouts"] == 0
+ assert data["completed_rollouts"] == 0
+
+ # Start a rollout
+ await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="count-1"),
+ )
+ await asyncio.sleep(0.05)
+
+ health = await client.get("/health")
+ data = health.json()
+ assert data["active_rollouts"] >= 1
+
+ # Wait for completion
+ await asyncio.sleep(0.5)
+
+ health = await client.get("/health")
+ data = health.json()
+ assert data["completed_rollouts"] >= 1
+
+ def test_health_includes_agent_loop_name(self) -> None:
+ """Health response should contain the agent loop name."""
+ app = create_app(SimpleAgentLoop())
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ data = resp.json()
+ assert data["agent_loop"] == "simple_agent"
+
+ def test_health_returns_status_healthy(self) -> None:
+ """Health endpoint should return status 'healthy'."""
+ app = create_app(SimpleAgentLoop())
+ with TestClient(app) as client:
+ resp = client.get("/health")
+ assert resp.status_code == 200
+ assert resp.json()["status"] == "healthy"
+
+
+# =============================================================================
+# create_app configuration variations
+# =============================================================================
+
+
+class TestCreateAppConfig:
+ """Tests for various create_app() parameter combinations."""
+
+ def test_custom_settings_used(self) -> None:
+ """Custom RolloutSettings should be respected."""
+ from osmosis_ai.rollout.config.settings import (
+ RolloutServerSettings,
+ RolloutSettings,
+ )
+
+ settings = RolloutSettings(
+ server=RolloutServerSettings(max_concurrent_rollouts=42)
+ )
+ app = create_app(SimpleAgentLoop(), settings=settings)
+ assert app is not None
+
+ def test_custom_record_ttl(self) -> None:
+ """Custom record_ttl_seconds should be accepted."""
+ app = create_app(SimpleAgentLoop(), record_ttl_seconds=120.0)
+ assert app is not None
+
+ def test_debug_dir_parameter(self) -> None:
+ """debug_dir parameter should be accepted."""
+ app = create_app(SimpleAgentLoop(), debug_dir="/tmp/debug-test")
+ assert app is not None
+
+ def test_app_title_includes_agent_name(self) -> None:
+ """The FastAPI app title should include the agent loop name."""
+ app = create_app(SimpleAgentLoop())
+ assert "simple_agent" in app.title
+
+ def test_default_settings_loaded_when_none(self) -> None:
+ """When settings=None, default settings should be loaded via get_settings()."""
+ app = create_app(SimpleAgentLoop())
+ assert app is not None
+ assert "simple_agent" in app.title
+
+
+# =============================================================================
+# Debug dir in background rollout
+# =============================================================================
+
+
+class TestDebugDirInRollout:
+ """Tests that debug_dir is passed through to the RolloutContext."""
+
+ async def test_debug_dir_propagated_to_context(self, tmp_path: Any) -> None:
+ """The debug_dir should be passed to RolloutContext during rollout."""
+ debug_dir = str(tmp_path / "debug_out")
+ observed_debug_dir: str | None = "UNSET"
+
+ class InspectingAgent(RolloutAgentLoop):
+ name = "inspecting_agent"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ nonlocal observed_debug_dir
+ observed_debug_dir = ctx._debug_dir
+ return ctx.complete([])
+
+ app = create_app(InspectingAgent(), debug_dir=debug_dir)
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="debug-prop"),
+ )
+ await asyncio.sleep(0.2)
+
+ assert observed_debug_dir == debug_dir
+
+ async def test_no_debug_dir_means_none_in_context(self) -> None:
+ """Without debug_dir, RolloutContext._debug_dir should be None."""
+ observed_debug_dir: str | None = "UNSET"
+
+ class InspectingAgent(RolloutAgentLoop):
+ name = "inspect_no_debug"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ nonlocal observed_debug_dir
+ observed_debug_dir = ctx._debug_dir
+ return ctx.complete([])
+
+ app = create_app(InspectingAgent())
+ mock_client = _mock_llm_client()
+ with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
+ MockCls.return_value = mock_client
+
+ transport = ASGITransport(app=app)
+ async with AsyncClient(
+ transport=transport, base_url="http://test"
+ ) as client:
+ await client.post(
+ "/v1/rollout/init",
+ json=_make_rollout_payload(rollout_id="no-debug"),
+ )
+ await asyncio.sleep(0.2)
+
+ assert observed_debug_dir is None
+
+
+# =============================================================================
+# ImportError guard
+# =============================================================================
+
+
+class TestImportGuard:
+ """Test that create_app() raises ImportError when FastAPI is absent."""
+
+ def test_raises_import_error_when_fastapi_missing(self) -> None:
+ """create_app() should raise ImportError with a helpful message."""
+ with patch("osmosis_ai.rollout.server.app.FASTAPI_AVAILABLE", False):
+ with pytest.raises(ImportError, match="FastAPI is required"):
+ create_app(SimpleAgentLoop())
diff --git a/tests/test_rollout_testing.py b/tests/test_rollout_testing.py
index daff4d58..ec25353f 100644
--- a/tests/test_rollout_testing.py
+++ b/tests/test_rollout_testing.py
@@ -16,13 +16,17 @@
from __future__ import annotations
+from unittest.mock import patch
+
import pytest
from osmosis_ai.rollout.testing import (
RolloutCompletionTracker,
+ _should_use_tools,
create_mock_trainer_app,
fake_prompt_token_ids,
fake_token_ids,
+ patch_httpx_for_mock_trainer,
)
# Check if FastAPI is available for testing
@@ -59,6 +63,20 @@ def test_fake_token_ids_length() -> None:
assert len(result) == len(text)
+def test_fake_token_ids_deterministic() -> None:
+ """Verify fake_token_ids returns the same result for the same input."""
+ result1 = fake_token_ids("test string")
+ result2 = fake_token_ids("test string")
+ assert result1 == result2
+
+
+def test_fake_token_ids_always_starts_at_zero() -> None:
+ """Verify fake_token_ids always starts the sequence at 0."""
+ result = fake_token_ids("abc")
+ assert result[0] == 0
+ assert result[-1] == 2
+
+
# =============================================================================
# fake_prompt_token_ids Tests
# =============================================================================
@@ -88,6 +106,24 @@ def test_fake_prompt_token_ids_multiple() -> None:
assert len(result) == 30 # 10 * 3
+def test_fake_prompt_token_ids_returns_sequential_ids() -> None:
+ """Verify fake_prompt_token_ids returns sequential integers starting at 0."""
+ messages = [
+ {"role": "user", "content": "Hi"},
+ {"role": "assistant", "content": "Hello"},
+ ]
+ result = fake_prompt_token_ids(messages)
+ assert result == list(range(20))
+
+
+def test_fake_prompt_token_ids_scales_linearly() -> None:
+ """Verify that each additional message adds exactly 10 tokens."""
+ for n in range(1, 6):
+ messages = [{"role": "user", "content": f"msg {i}"} for i in range(n)]
+ result = fake_prompt_token_ids(messages)
+ assert len(result) == 10 * n
+
+
# =============================================================================
# RolloutCompletionTracker Tests
# =============================================================================
@@ -144,6 +180,143 @@ def test_tracker_wait_success() -> None:
assert result is True
+def test_tracker_clear_after_multiple_records() -> None:
+ """Verify tracker clear removes all recorded responses and resets the event."""
+ tracker = RolloutCompletionTracker()
+ tracker.record({"rollout_id": "r1"})
+ tracker.record({"rollout_id": "r2"})
+ tracker.record({"rollout_id": "r3"})
+
+ assert len(tracker.responses) == 3
+ assert tracker.event.is_set()
+
+ tracker.clear()
+
+ assert tracker.responses == []
+ assert not tracker.event.is_set()
+ # After clearing, wait should timeout
+ assert tracker.wait(timeout=0.01) is False
+
+
+def test_tracker_record_after_clear() -> None:
+ """Verify tracker can record new responses after being cleared."""
+ tracker = RolloutCompletionTracker()
+ tracker.record({"rollout_id": "old"})
+ tracker.clear()
+ tracker.record({"rollout_id": "new"})
+
+ assert len(tracker.responses) == 1
+ assert tracker.responses[0]["rollout_id"] == "new"
+ assert tracker.event.is_set()
+
+
+# =============================================================================
+# _should_use_tools Tests
+# =============================================================================
+
+
+def test_should_use_tools_returns_true_for_calculate_keyword() -> None:
+ """Verify _should_use_tools detects 'calculate' in user message."""
+ msg = {"role": "user", "content": "Please calculate 5 + 3"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_add_keyword() -> None:
+ """Verify _should_use_tools detects 'add' in user message."""
+ msg = {"role": "user", "content": "Can you add these numbers?"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_sum_keyword() -> None:
+ """Verify _should_use_tools detects 'sum' in user message."""
+ msg = {"role": "user", "content": "What is the sum of 10 and 20?"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_plus_keyword() -> None:
+ """Verify _should_use_tools detects 'plus' in user message."""
+ msg = {"role": "user", "content": "5 plus 3 equals what?"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_subtract_keyword() -> None:
+ """Verify _should_use_tools detects 'subtract' in user message."""
+ msg = {"role": "user", "content": "subtract 3 from 10"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_multiply_keyword() -> None:
+ """Verify _should_use_tools detects 'multiply' in user message."""
+ msg = {"role": "user", "content": "multiply 4 by 5"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_true_for_divide_keyword() -> None:
+ """Verify _should_use_tools detects 'divide' in user message."""
+ msg = {"role": "user", "content": "divide 20 by 4"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_is_case_insensitive() -> None:
+ """Verify _should_use_tools matches keywords regardless of case."""
+ msg = {"role": "user", "content": "CALCULATE the total"}
+ assert _should_use_tools(msg) is True
+
+
+def test_should_use_tools_returns_false_for_non_calculator_message() -> None:
+ """Verify _should_use_tools returns False for generic messages."""
+ msg = {"role": "user", "content": "Hello, how are you?"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_assistant_role() -> None:
+ """Verify _should_use_tools returns False when role is not 'user'."""
+ msg = {"role": "assistant", "content": "I will calculate that for you"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_system_role() -> None:
+ """Verify _should_use_tools returns False for system messages even with keywords."""
+ msg = {"role": "system", "content": "You can add numbers"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_tool_role() -> None:
+ """Verify _should_use_tools returns False for tool messages."""
+ msg = {"role": "tool", "content": "The sum is 8"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_non_string_content() -> None:
+ """Verify _should_use_tools returns False when content is not a string."""
+ msg = {"role": "user", "content": 12345}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_none_content() -> None:
+ """Verify _should_use_tools returns False when content is None."""
+ msg = {"role": "user", "content": None}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_missing_content() -> None:
+ """Verify _should_use_tools returns False when content key is missing."""
+ msg = {"role": "user"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_missing_role() -> None:
+ """Verify _should_use_tools returns False when role key is missing."""
+ msg = {"content": "calculate 5 + 3"}
+ assert _should_use_tools(msg) is False
+
+
+def test_should_use_tools_returns_false_for_list_content() -> None:
+ """Verify _should_use_tools returns False when content is a list (multimodal)."""
+ msg = {"role": "user", "content": [{"type": "text", "text": "calculate 5+3"}]}
+ assert _should_use_tools(msg) is False
+
+
# =============================================================================
# create_mock_trainer_app Tests
# =============================================================================
@@ -359,3 +532,301 @@ def custom_generator(message):
data = response.json()
message = data["choices"][0]["message"]
assert "tool_calls" not in message or message.get("tool_calls") is None
+
+
+@requires_fastapi
+def test_mock_trainer_completions_with_empty_messages() -> None:
+ """Verify mock trainer handles empty messages list gracefully."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "test-empty",
+ "messages": [],
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ # When messages list is empty, it uses default placeholder message
+ assert data["choices"][0]["message"]["role"] == "assistant"
+ assert data["choices"][0]["message"]["content"] == "OK."
+
+
+@requires_fastapi
+def test_mock_trainer_completions_token_ids_match_response_length() -> None:
+ """Verify token IDs length matches response content length."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "test-tokens",
+ "messages": [{"role": "user", "content": "Hello"}],
+ },
+ )
+ data = response.json()
+ content = data["choices"][0]["message"]["content"]
+ assert len(data["token_ids"]) == len(content)
+ assert len(data["logprobs"]) == len(data["token_ids"])
+
+
+@requires_fastapi
+def test_mock_trainer_completions_usage_totals_are_consistent() -> None:
+ """Verify that usage total_tokens equals prompt_tokens + completion_tokens."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "test-usage",
+ "messages": [
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": "Hi"},
+ {"role": "user", "content": "Tell me more"},
+ ],
+ },
+ )
+ data = response.json()
+ usage = data["usage"]
+ assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+
+
+@requires_fastapi
+def test_mock_trainer_completions_tool_call_response_has_correct_structure() -> None:
+ """Verify tool call response includes content alongside tool_calls."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "test-tool-struct",
+ "messages": [{"role": "user", "content": "Please add 5 and 3"}],
+ },
+ )
+ data = response.json()
+ message = data["choices"][0]["message"]
+
+ # Tool call responses should have both content and tool_calls
+ assert message["content"] == "I'll help you with that calculation."
+ assert len(message["tool_calls"]) == 1
+ tool_call = message["tool_calls"][0]
+ assert tool_call["id"] == "call_mock_add"
+ assert tool_call["type"] == "function"
+ assert tool_call["function"]["name"] == "add"
+ assert tool_call["function"]["arguments"] == '{"a": 5, "b": 3}'
+
+
+@requires_fastapi
+def test_mock_trainer_without_tracker_still_stores_rollouts() -> None:
+ """Verify completed rollouts are stored even without a tracker."""
+ app = create_mock_trainer_app(tracker=None)
+ client = TestClient(app)
+
+ client.post(
+ "/v1/rollout/completed",
+ json={
+ "rollout_id": "no-tracker-123",
+ "status": "COMPLETED",
+ "final_messages": [{"role": "assistant", "content": "Done"}],
+ "finish_reason": "stop",
+ },
+ )
+
+ response = client.get("/v1/rollout/completed/no-tracker-123")
+ assert response.status_code == 200
+ data = response.json()
+ assert data["rollout_id"] == "no-tracker-123"
+
+
+@requires_fastapi
+def test_mock_trainer_completions_model_field_echoed() -> None:
+ """Verify the model field from request is echoed in response."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "test-model",
+ "messages": [{"role": "user", "content": "Hi"}],
+ "model": "custom-model-v1",
+ },
+ )
+ data = response.json()
+ assert data["model"] == "custom-model-v1"
+
+
+@requires_fastapi
+def test_mock_trainer_completions_rollout_id_echoed_as_response_id() -> None:
+ """Verify rollout_id is used as the response id."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "rollout_id": "unique-id-xyz",
+ "messages": [{"role": "user", "content": "Hi"}],
+ },
+ )
+ data = response.json()
+ assert data["id"] == "unique-id-xyz"
+
+
+@requires_fastapi
+def test_mock_trainer_completed_rollout_overwrites_previous() -> None:
+ """Verify completing the same rollout_id overwrites the previous entry."""
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+
+ # First completion
+ client.post(
+ "/v1/rollout/completed",
+ json={
+ "rollout_id": "overwrite-test",
+ "status": "COMPLETED",
+ "final_messages": [{"role": "assistant", "content": "First"}],
+ "finish_reason": "stop",
+ },
+ )
+
+ # Second completion with same rollout_id
+ client.post(
+ "/v1/rollout/completed",
+ json={
+ "rollout_id": "overwrite-test",
+ "status": "ERROR",
+ "final_messages": [{"role": "assistant", "content": "Second"}],
+ "finish_reason": "error",
+ "error_message": "something failed",
+ },
+ )
+
+ response = client.get("/v1/rollout/completed/overwrite-test")
+ data = response.json()
+ assert data["status"] == "ERROR"
+ assert data["finish_reason"] == "error"
+
+
+def test_create_mock_trainer_app_raises_import_error_when_fastapi_missing() -> None:
+ """Verify create_mock_trainer_app raises ImportError when FastAPI is not installed."""
+ import builtins
+
+ original_import = builtins.__import__
+
+ def mock_import(name, *args, **kwargs):
+ if name == "fastapi":
+ raise ImportError("No module named 'fastapi'")
+ return original_import(name, *args, **kwargs)
+
+ with patch.object(builtins, "__import__", side_effect=mock_import):
+ with pytest.raises(ImportError, match="FastAPI is required"):
+ create_mock_trainer_app()
+
+
+# =============================================================================
+# patch_httpx_for_mock_trainer Tests
+# =============================================================================
+
+
+@requires_fastapi
+async def test_patch_httpx_routes_completions_to_mock_trainer(monkeypatch) -> None:
+ """Verify patch_httpx_for_mock_trainer routes /v1/chat/completions to the mock."""
+ import httpx
+
+ tracker = RolloutCompletionTracker()
+ app = create_mock_trainer_app(tracker=tracker)
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+
+ async with httpx.AsyncClient() as http_client:
+ response = await http_client.post(
+ "http://fake-server:8080/v1/chat/completions",
+ json={
+ "rollout_id": "patched-test",
+ "messages": [{"role": "user", "content": "Hello"}],
+ },
+ )
+
+ assert response.status_code == 200
+ data = response.json()
+ assert data["id"] == "patched-test"
+ assert data["choices"][0]["message"]["role"] == "assistant"
+
+
+@requires_fastapi
+async def test_patch_httpx_routes_rollout_completed_to_mock_trainer(
+ monkeypatch,
+) -> None:
+ """Verify patch_httpx_for_mock_trainer routes /v1/rollout/completed to the mock."""
+ import httpx
+
+ tracker = RolloutCompletionTracker()
+ app = create_mock_trainer_app(tracker=tracker)
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+
+ async with httpx.AsyncClient() as http_client:
+ response = await http_client.post(
+ "http://fake-server:8080/v1/rollout/completed",
+ json={
+ "rollout_id": "completed-test",
+ "status": "COMPLETED",
+ "final_messages": [],
+ "finish_reason": "stop",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json()["status"] == "ok"
+ assert len(tracker.responses) == 1
+ assert tracker.responses[0]["rollout_id"] == "completed-test"
+
+
+@requires_fastapi
+async def test_patch_httpx_does_not_intercept_unrelated_urls(monkeypatch) -> None:
+ """Verify patch_httpx_for_mock_trainer does not intercept non-rollout URLs."""
+ import httpx
+
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+
+ with pytest.raises((httpx.ConnectError, httpx.ConnectTimeout)):
+ async with httpx.AsyncClient() as http_client:
+ await http_client.post(
+ "http://127.0.0.1:1/v1/some/other/endpoint",
+ json={"data": "test"},
+ timeout=0.1,
+ )
+
+
+@requires_fastapi
+async def test_patch_httpx_completions_response_is_valid_httpx_response(
+ monkeypatch,
+) -> None:
+ """Verify the patched response is a proper httpx.Response object."""
+ import httpx
+
+ app = create_mock_trainer_app()
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+
+ async with httpx.AsyncClient() as http_client:
+ response = await http_client.post(
+ "http://any-host/v1/chat/completions",
+ json={
+ "rollout_id": "response-type-test",
+ "messages": [{"role": "user", "content": "Hello"}],
+ },
+ )
+
+ assert isinstance(response, httpx.Response)
+ assert response.status_code == 200
+ assert response.request.method == "POST"
+ assert "v1/chat/completions" in str(response.request.url)
diff --git a/tests/test_utils_decorators.py b/tests/test_utils_decorators.py
new file mode 100644
index 00000000..8f76f33c
--- /dev/null
+++ b/tests/test_utils_decorators.py
@@ -0,0 +1,946 @@
+"""Tests for osmosis_ai.utils decorator validation, async wrapping, and type checking.
+
+Covers @osmosis_reward and @osmosis_rubric decorators with both valid and invalid
+function signatures, async function support, return type enforcement, and the
+_is_str_annotation helper.
+
+NOTE: This file intentionally does NOT use ``from __future__ import annotations``
+because osmosis_reward checks annotation identity (``annotation is not str``),
+which breaks when annotations are stringified by PEP 563.
+"""
+
+import sys
+
+import pytest
+
+from osmosis_ai.utils import _is_str_annotation, osmosis_reward, osmosis_rubric
+
+# =============================================================================
+# _is_str_annotation Tests
+# =============================================================================
+
+
+class TestIsStrAnnotation:
+ """Tests for the _is_str_annotation helper used by osmosis_rubric."""
+
+ def test_str_type_is_recognized(self) -> None:
+ """The builtin str type should be recognized."""
+ assert _is_str_annotation(str) is True
+
+ def test_str_string_literal_is_recognized(self) -> None:
+ """String literal 'str' should be recognized."""
+ assert _is_str_annotation("str") is True
+
+ def test_builtins_str_string_is_recognized(self) -> None:
+ """String literal 'builtins.str' should be recognized."""
+ assert _is_str_annotation("builtins.str") is True
+
+ def test_empty_annotation_rejected(self) -> None:
+ """inspect.Parameter.empty (no annotation) should not match str."""
+ import inspect
+
+ assert _is_str_annotation(inspect.Parameter.empty) is False
+
+ def test_non_str_type_rejected(self) -> None:
+ """Other types (int, dict, etc.) should be rejected."""
+ assert _is_str_annotation(int) is False
+ assert _is_str_annotation(dict) is False
+ assert _is_str_annotation(float) is False
+
+ def test_non_str_string_rejected(self) -> None:
+ """Arbitrary strings that are not 'str' or 'builtins.str' are rejected."""
+ assert _is_str_annotation("int") is False
+ assert _is_str_annotation("something") is False
+
+ def test_str_subclass_is_recognized(self) -> None:
+ """A subclass of str should be recognized via issubclass."""
+
+ class MyStr(str):
+ pass
+
+ assert _is_str_annotation(MyStr) is True
+
+ def test_none_type_rejected(self) -> None:
+ """NoneType should be rejected."""
+ assert _is_str_annotation(type(None)) is False
+
+
+# =============================================================================
+# @osmosis_reward - Valid Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRewardValidClassic:
+ """Tests for valid classic-style reward functions."""
+
+ def test_minimal_classic_signature(self) -> None:
+ """Classic signature with exactly 3 required params is accepted."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict = None, # noqa: RUF013
+ ) -> float:
+ return 1.0
+
+ assert reward("answer", "truth", None) == 1.0
+
+ def test_classic_with_kwargs(self) -> None:
+ """Classic signature with **kwargs is accepted and kwargs are passed through."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ **kwargs,
+ ) -> float:
+ return 1.0
+
+ # data_source should be passed through when **kwargs is present
+ assert reward("a", "b", None, data_source="test") == 1.0
+
+ def test_classic_with_optional_dict_annotation(self) -> None:
+ """Classic signature with Optional[dict] for extra_info is accepted."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 0.5
+
+ assert reward("a", "b") == 0.5
+
+ def test_classic_with_union_dict_none_annotation(self) -> None:
+ """Classic signature with Union[dict, None] for extra_info is accepted."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 0.5
+
+ assert reward("a", "b") == 0.5
+
+ @pytest.mark.skipif(
+ sys.version_info < (3, 10), reason="PEP 604 syntax requires Python 3.10+"
+ )
+ def test_classic_with_pep604_union(self) -> None:
+ """Classic signature with dict | None (PEP 604) for extra_info is accepted."""
+ # Use exec to avoid SyntaxError on Python < 3.10
+ code = """
+def _reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+) -> float:
+ return 0.5
+"""
+ local_ns: dict = {}
+ exec(code, local_ns)
+ wrapped = osmosis_reward(local_ns["_reward"])
+ assert wrapped("a", "b") == 0.5
+
+ def test_return_value_is_float(self) -> None:
+ """Classic reward function should enforce float return."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 0.75
+
+ result = reward("a", "b")
+ assert isinstance(result, float)
+ assert result == 0.75
+
+
+# =============================================================================
+# @osmosis_reward - Invalid Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRewardInvalidClassic:
+ """Tests for classic-style reward functions that should be rejected at decoration time."""
+
+ def test_too_few_parameters(self) -> None:
+ """Classic signature with fewer than 3 params should raise TypeError."""
+ with pytest.raises(TypeError, match="must have at least 3 parameters"):
+
+ @osmosis_reward
+ def reward(solution_str: str, ground_truth: str) -> float:
+ return 1.0
+
+ def test_solution_str_wrong_type(self) -> None:
+ """First parameter annotated as non-str should raise TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as str"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: int,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_ground_truth_wrong_name(self) -> None:
+ """Second parameter with wrong name should raise TypeError."""
+ with pytest.raises(TypeError, match="must be named 'ground_truth'"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ wrong_name: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_ground_truth_wrong_type(self) -> None:
+ """Second parameter annotated as non-str should raise TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as str"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: int,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_extra_info_wrong_name(self) -> None:
+ """Third parameter with wrong name should raise TypeError."""
+ with pytest.raises(TypeError, match="must be named 'extra_info'"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ wrong_name: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_extra_info_wrong_type(self) -> None:
+ """Third parameter annotated as non-dict should raise TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as dict"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: str = None, # noqa: RUF013
+ ) -> float:
+ return 1.0
+
+ def test_extra_info_no_default(self) -> None:
+ """Third parameter without default value should raise TypeError."""
+ with pytest.raises(TypeError, match="must have a default value"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ def test_extra_info_union_with_wrong_type(self) -> None:
+ """Third parameter with Union[str, None] should raise TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as dict"):
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: str | None = None,
+ ) -> float:
+ return 1.0
+
+
+# =============================================================================
+# @osmosis_reward - Non-Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRewardNonClassic:
+ """Tests for non-classic reward functions (first param is NOT solution_str)."""
+
+ def test_arbitrary_signature_accepted(self) -> None:
+ """Non-classic signatures are accepted without validation."""
+
+ @osmosis_reward
+ def custom_reward(messages: list, metadata: dict) -> float:
+ return 0.5
+
+ assert custom_reward([], {}) == 0.5
+
+ def test_no_params_accepted(self) -> None:
+ """Zero-param functions are accepted in non-classic mode."""
+
+ @osmosis_reward
+ def no_params() -> float:
+ return 1.0
+
+ assert no_params() == 1.0
+
+ def test_non_classic_no_return_type_enforcement(self) -> None:
+ """Non-classic functions have no return type enforcement."""
+
+ @osmosis_reward
+ def reward(x: int) -> str:
+ return "not a float"
+
+ # No TypeError is raised even though return is not float
+ assert reward(42) == "not a float"
+
+
+# =============================================================================
+# @osmosis_reward - Return Type Enforcement (Classic)
+# =============================================================================
+
+
+class TestOsmosisRewardReturnType:
+ """Tests for classic return type enforcement at call time."""
+
+ def test_non_float_return_raises_typeerror(self) -> None:
+ """Classic reward returning non-float raises TypeError at call time."""
+
+ @osmosis_reward
+ def bad_return(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return "not a float" # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ bad_return("a", "b")
+
+ def test_int_return_raises_typeerror(self) -> None:
+ """Classic reward returning int (not float) raises TypeError."""
+
+ @osmosis_reward
+ def int_return(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1 # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ int_return("a", "b")
+
+
+# =============================================================================
+# @osmosis_reward - data_source Handling
+# =============================================================================
+
+
+class TestOsmosisRewardDataSource:
+ """Tests for the data_source kwarg stripping behavior."""
+
+ def test_data_source_stripped_without_kwargs(self) -> None:
+ """When function has no **kwargs and no data_source param, data_source is stripped."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ # This should not raise even though data_source is passed
+ assert reward("a", "b", data_source="test") == 1.0
+
+ def test_data_source_passed_when_kwargs_present(self) -> None:
+ """When function has **kwargs, data_source is NOT stripped."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ **kwargs,
+ ) -> float:
+ return 1.0 if "data_source" in kwargs else 0.0
+
+ assert reward("a", "b", data_source="test") == 1.0
+
+ def test_data_source_passed_when_explicit_param(self) -> None:
+ """When function has explicit data_source param, it is NOT stripped."""
+
+ @osmosis_reward
+ def reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ data_source: str = None, # noqa: RUF013
+ ) -> float:
+ return 1.0 if data_source == "test" else 0.0
+
+ assert reward("a", "b", data_source="test") == 1.0
+
+
+# =============================================================================
+# @osmosis_reward - Async Support
+# =============================================================================
+
+
+class TestOsmosisRewardAsync:
+ """Tests for async reward function wrapping."""
+
+ async def test_async_classic_reward(self) -> None:
+ """Async classic reward function is properly wrapped and awaitable."""
+
+ @osmosis_reward
+ async def async_reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 0.9
+
+ result = await async_reward("a", "b")
+ assert result == 0.9
+
+ async def test_async_return_type_enforcement(self) -> None:
+ """Async classic reward returning non-float raises TypeError."""
+
+ @osmosis_reward
+ async def bad_async(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return "bad" # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ await bad_async("a", "b")
+
+ async def test_async_data_source_stripped(self) -> None:
+ """Async reward function correctly strips data_source."""
+
+ @osmosis_reward
+ async def async_reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ # Should not raise even with data_source
+ result = await async_reward("a", "b", data_source="x")
+ assert result == 1.0
+
+ async def test_async_non_classic(self) -> None:
+ """Async non-classic reward is properly wrapped."""
+
+ @osmosis_reward
+ async def async_custom(x: int) -> int:
+ return x * 2
+
+ result = await async_custom(5)
+ assert result == 10
+
+
+# =============================================================================
+# @osmosis_reward - functools.wraps Preservation
+# =============================================================================
+
+
+class TestOsmosisRewardWraps:
+ """Tests that functools.wraps preserves function metadata."""
+
+ def test_function_name_preserved(self) -> None:
+ """Wrapped function retains its original __name__."""
+
+ @osmosis_reward
+ def my_reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ assert my_reward.__name__ == "my_reward"
+
+ def test_docstring_preserved(self) -> None:
+ """Wrapped function retains its docstring."""
+
+ @osmosis_reward
+ def documented_reward(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict | None = None,
+ ) -> float:
+ """This is my reward function."""
+ return 1.0
+
+ assert documented_reward.__doc__ == "This is my reward function."
+
+
+# =============================================================================
+# @osmosis_rubric - Valid Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRubricValidClassic:
+ """Tests for valid classic-style rubric functions."""
+
+ def test_minimal_classic_signature(self) -> None:
+ """Classic signature with exactly 3 required params is accepted."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ result = rubric("answer", "truth", {})
+ assert result == 1.0
+
+ def test_classic_with_kwargs(self) -> None:
+ """Classic signature with **kwargs is accepted."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ **kwargs,
+ ) -> float:
+ return 1.0
+
+ assert rubric("a", "b", {}) == 1.0
+
+ def test_ground_truth_optional_str(self) -> None:
+ """ground_truth annotated as Optional[str] is accepted."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str | None,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ assert rubric("a", None, {}) == 1.0
+
+ def test_ground_truth_union_str_none(self) -> None:
+ """ground_truth annotated as Union[str, None] is accepted."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str | None,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ assert rubric("a", None, {}) == 1.0
+
+
+# =============================================================================
+# @osmosis_rubric - Invalid Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRubricInvalidClassic:
+ """Tests for classic-style rubric functions that should be rejected."""
+
+ def test_too_few_parameters(self) -> None:
+ """Classic signature with fewer than 3 params should raise TypeError."""
+ with pytest.raises(TypeError, match="must have at least 3 parameters"):
+
+ @osmosis_rubric
+ def rubric(solution_str: str, ground_truth: str) -> float:
+ return 1.0
+
+ def test_solution_str_wrong_type(self) -> None:
+ """First parameter annotated as non-str should raise TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as str"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: int,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ def test_solution_str_with_default_rejected(self) -> None:
+ """First parameter with a default value should raise TypeError."""
+ with pytest.raises(TypeError, match="cannot have a default value"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str = "default",
+ ground_truth: str = "",
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_ground_truth_wrong_name(self) -> None:
+ """Second parameter with wrong name should raise TypeError."""
+ with pytest.raises(TypeError, match="must be named 'ground_truth'"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ wrong_name: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ def test_ground_truth_wrong_type(self) -> None:
+ """Second parameter annotated as non-str (and not Optional[str]) raises TypeError."""
+ with pytest.raises(TypeError, match="must be annotated as str"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: int,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ def test_ground_truth_with_default_rejected(self) -> None:
+ """Second parameter with a default value should raise TypeError."""
+ with pytest.raises(TypeError, match="cannot have a default value"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str = "default",
+ extra_info: dict | None = None,
+ ) -> float:
+ return 1.0
+
+ def test_extra_info_wrong_name(self) -> None:
+ """Third parameter with wrong name should raise TypeError."""
+ with pytest.raises(TypeError, match="must be named 'extra_info'"):
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ wrong_name: dict,
+ ) -> float:
+ return 1.0
+
+
+# =============================================================================
+# @osmosis_rubric - Runtime Argument Validation (Classic)
+# =============================================================================
+
+
+class TestOsmosisRubricRuntimeValidation:
+ """Tests for runtime argument validation in classic rubric functions."""
+
+ def test_solution_str_must_be_string(self) -> None:
+ """Passing a non-string for solution_str raises TypeError at call time."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ with pytest.raises(TypeError, match="must be a string"):
+ rubric(123, "truth", {}) # type: ignore[arg-type]
+
+ def test_ground_truth_must_be_string_or_none(self) -> None:
+ """Passing non-string/non-None for ground_truth raises TypeError."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str | None,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ with pytest.raises(TypeError, match="must be a string or None"):
+ rubric("answer", 123, {}) # type: ignore[arg-type]
+
+ def test_ground_truth_none_is_accepted(self) -> None:
+ """None is a valid value for ground_truth."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str | None,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ assert rubric("a", None, {}) == 1.0
+
+ def test_extra_info_required(self) -> None:
+ """extra_info is required for classic rubric functions."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ with pytest.raises(TypeError, match="'extra_info' argument is required"):
+ rubric("a", "b") # type: ignore[call-arg]
+
+ def test_solution_str_required(self) -> None:
+ """solution_str is required for classic rubric functions."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ with pytest.raises(TypeError):
+ rubric() # type: ignore[call-arg]
+
+
+# =============================================================================
+# @osmosis_rubric - Return Type Enforcement (Classic)
+# =============================================================================
+
+
+class TestOsmosisRubricReturnType:
+ """Tests for classic rubric return type enforcement."""
+
+ def test_non_float_return_raises_typeerror(self) -> None:
+ """Classic rubric returning non-float raises TypeError."""
+
+ @osmosis_rubric
+ def bad_return(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return "not a float" # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ bad_return("a", "b", {})
+
+ def test_int_return_raises_typeerror(self) -> None:
+ """Classic rubric returning int raises TypeError."""
+
+ @osmosis_rubric
+ def int_return(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1 # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ int_return("a", "b", {})
+
+
+# =============================================================================
+# @osmosis_rubric - Non-Classic Signatures
+# =============================================================================
+
+
+class TestOsmosisRubricNonClassic:
+ """Tests for non-classic rubric functions."""
+
+ def test_arbitrary_signature_accepted(self) -> None:
+ """Non-classic signatures are accepted without validation."""
+
+ @osmosis_rubric
+ def custom(messages: list, context: dict) -> float:
+ return 0.5
+
+ assert custom([], {}) == 0.5
+
+ def test_no_params_accepted(self) -> None:
+ """Zero-param functions are accepted in non-classic mode."""
+
+ @osmosis_rubric
+ def no_params() -> float:
+ return 1.0
+
+ assert no_params() == 1.0
+
+ def test_non_classic_no_return_type_enforcement(self) -> None:
+ """Non-classic functions have no return type enforcement."""
+
+ @osmosis_rubric
+ def rubric(x: int) -> str:
+ return "string"
+
+ assert rubric(1) == "string"
+
+ def test_non_classic_no_runtime_arg_validation(self) -> None:
+ """Non-classic functions skip runtime argument validation."""
+
+ @osmosis_rubric
+ def rubric(x: int) -> int:
+ return x * 2
+
+ # No TypeError even though int is passed
+ assert rubric(5) == 10
+
+
+# =============================================================================
+# @osmosis_rubric - data_source Handling
+# =============================================================================
+
+
+class TestOsmosisRubricDataSource:
+ """Tests for the data_source kwarg stripping behavior in rubric."""
+
+ def test_data_source_stripped_without_kwargs(self) -> None:
+ """When function has no **kwargs, data_source is stripped."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ assert rubric("a", "b", {}, data_source="test") == 1.0
+
+ def test_data_source_passed_when_kwargs_present(self) -> None:
+ """When function has **kwargs, data_source is NOT stripped."""
+
+ @osmosis_rubric
+ def rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ **kwargs,
+ ) -> float:
+ return 1.0 if "data_source" in kwargs else 0.0
+
+ assert rubric("a", "b", {}, data_source="test") == 1.0
+
+
+# =============================================================================
+# @osmosis_rubric - Async Support
+# =============================================================================
+
+
+class TestOsmosisRubricAsync:
+ """Tests for async rubric function wrapping."""
+
+ async def test_async_classic_rubric(self) -> None:
+ """Async classic rubric function is properly wrapped and awaitable."""
+
+ @osmosis_rubric
+ async def async_rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 0.8
+
+ result = await async_rubric("a", "b", {})
+ assert result == 0.8
+
+ async def test_async_return_type_enforcement(self) -> None:
+ """Async classic rubric returning non-float raises TypeError."""
+
+ @osmosis_rubric
+ async def bad_async(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return "bad" # type: ignore[return-value]
+
+ with pytest.raises(TypeError, match="must return a float"):
+ await bad_async("a", "b", {})
+
+ async def test_async_runtime_validation(self) -> None:
+ """Async rubric validates arguments at call time."""
+
+ @osmosis_rubric
+ async def async_rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ with pytest.raises(TypeError, match="must be a string"):
+ await async_rubric(42, "b", {}) # type: ignore[arg-type]
+
+ async def test_async_data_source_stripped(self) -> None:
+ """Async rubric correctly strips data_source."""
+
+ @osmosis_rubric
+ async def async_rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ result = await async_rubric("a", "b", {}, data_source="x")
+ assert result == 1.0
+
+ async def test_async_non_classic(self) -> None:
+ """Async non-classic rubric is properly wrapped."""
+
+ @osmosis_rubric
+ async def async_custom(x: int) -> int:
+ return x * 2
+
+ result = await async_custom(5)
+ assert result == 10
+
+
+# =============================================================================
+# @osmosis_rubric - functools.wraps Preservation
+# =============================================================================
+
+
+class TestOsmosisRubricWraps:
+ """Tests that functools.wraps preserves function metadata."""
+
+ def test_function_name_preserved(self) -> None:
+ """Wrapped rubric retains its original __name__."""
+
+ @osmosis_rubric
+ def my_rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ return 1.0
+
+ assert my_rubric.__name__ == "my_rubric"
+
+ def test_docstring_preserved(self) -> None:
+ """Wrapped rubric retains its docstring."""
+
+ @osmosis_rubric
+ def documented_rubric(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict,
+ ) -> float:
+ """Evaluate quality."""
+ return 1.0
+
+ assert documented_rubric.__doc__ == "Evaluate quality."
diff --git a/uv.lock b/uv.lock
index 9f6b88d5..74424eba 100644
--- a/uv.lock
+++ b/uv.lock
@@ -465,6 +465,124 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
]
+[[package]]
+name = "coverage"
+version = "7.13.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/24/56/95b7e30fa389756cb56630faa728da46a27b8c6eb46f9d557c68fff12b65/coverage-7.13.4.tar.gz", hash = "sha256:e5c8f6ed1e61a8b2dcdf31eb0b9bbf0130750ca79c1c49eb898e2ad86f5ccc91", size = 827239, upload-time = "2026-02-09T12:59:03.86Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/44/d4/7827d9ffa34d5d4d752eec907022aa417120936282fc488306f5da08c292/coverage-7.13.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0fc31c787a84f8cd6027eba44010517020e0d18487064cd3d8968941856d1415", size = 219152, upload-time = "2026-02-09T12:56:11.974Z" },
+ { url = "https://files.pythonhosted.org/packages/35/b0/d69df26607c64043292644dbb9dc54b0856fabaa2cbb1eeee3331cc9e280/coverage-7.13.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a32ebc02a1805adf637fc8dec324b5cdacd2e493515424f70ee33799573d661b", size = 219667, upload-time = "2026-02-09T12:56:13.33Z" },
+ { url = "https://files.pythonhosted.org/packages/82/a4/c1523f7c9e47b2271dbf8c2a097e7a1f89ef0d66f5840bb59b7e8814157b/coverage-7.13.4-cp310-cp310-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:e24f9156097ff9dc286f2f913df3a7f63c0e333dcafa3c196f2c18b4175ca09a", size = 246425, upload-time = "2026-02-09T12:56:14.552Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/02/aa7ec01d1a5023c4b680ab7257f9bfde9defe8fdddfe40be096ac19e8177/coverage-7.13.4-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8041b6c5bfdc03257666e9881d33b1abc88daccaf73f7b6340fb7946655cd10f", size = 248229, upload-time = "2026-02-09T12:56:16.31Z" },
+ { url = "https://files.pythonhosted.org/packages/35/98/85aba0aed5126d896162087ef3f0e789a225697245256fc6181b95f47207/coverage-7.13.4-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2a09cfa6a5862bc2fc6ca7c3def5b2926194a56b8ab78ffcf617d28911123012", size = 250106, upload-time = "2026-02-09T12:56:18.024Z" },
+ { url = "https://files.pythonhosted.org/packages/96/72/1db59bd67494bc162e3e4cd5fbc7edba2c7026b22f7c8ef1496d58c2b94c/coverage-7.13.4-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:296f8b0af861d3970c2a4d8c91d48eb4dd4771bcef9baedec6a9b515d7de3def", size = 252021, upload-time = "2026-02-09T12:56:19.272Z" },
+ { url = "https://files.pythonhosted.org/packages/9d/97/72899c59c7066961de6e3daa142d459d47d104956db43e057e034f015c8a/coverage-7.13.4-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e101609bcbbfb04605ea1027b10dc3735c094d12d40826a60f897b98b1c30256", size = 247114, upload-time = "2026-02-09T12:56:21.051Z" },
+ { url = "https://files.pythonhosted.org/packages/39/1f/f1885573b5970235e908da4389176936c8933e86cb316b9620aab1585fa2/coverage-7.13.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:aa3feb8db2e87ff5e6d00d7e1480ae241876286691265657b500886c98f38bda", size = 248143, upload-time = "2026-02-09T12:56:22.585Z" },
+ { url = "https://files.pythonhosted.org/packages/a8/cf/e80390c5b7480b722fa3e994f8202807799b85bc562aa4f1dde209fbb7be/coverage-7.13.4-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:4fc7fa81bbaf5a02801b65346c8b3e657f1d93763e58c0abdf7c992addd81a92", size = 246152, upload-time = "2026-02-09T12:56:23.748Z" },
+ { url = "https://files.pythonhosted.org/packages/44/bf/f89a8350d85572f95412debb0fb9bb4795b1d5b5232bd652923c759e787b/coverage-7.13.4-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:33901f604424145c6e9c2398684b92e176c0b12df77d52db81c20abd48c3794c", size = 249959, upload-time = "2026-02-09T12:56:25.209Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/6e/612a02aece8178c818df273e8d1642190c4875402ca2ba74514394b27aba/coverage-7.13.4-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:bb28c0f2cf2782508a40cec377935829d5fcc3ad9a3681375af4e84eb34b6b58", size = 246416, upload-time = "2026-02-09T12:56:26.475Z" },
+ { url = "https://files.pythonhosted.org/packages/cb/98/b5afc39af67c2fa6786b03c3a7091fc300947387ce8914b096db8a73d67a/coverage-7.13.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:9d107aff57a83222ddbd8d9ee705ede2af2cc926608b57abed8ef96b50b7e8f9", size = 247025, upload-time = "2026-02-09T12:56:27.727Z" },
+ { url = "https://files.pythonhosted.org/packages/51/30/2bba8ef0682d5bd210c38fe497e12a06c9f8d663f7025e9f5c2c31ce847d/coverage-7.13.4-cp310-cp310-win32.whl", hash = "sha256:a6f94a7d00eb18f1b6d403c91a88fd58cfc92d4b16080dfdb774afc8294469bf", size = 221758, upload-time = "2026-02-09T12:56:29.051Z" },
+ { url = "https://files.pythonhosted.org/packages/78/13/331f94934cf6c092b8ea59ff868eb587bc8fe0893f02c55bc6c0183a192e/coverage-7.13.4-cp310-cp310-win_amd64.whl", hash = "sha256:2cb0f1e000ebc419632bbe04366a8990b6e32c4e0b51543a6484ffe15eaeda95", size = 222693, upload-time = "2026-02-09T12:56:30.366Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/ad/b59e5b451cf7172b8d1043dc0fa718f23aab379bc1521ee13d4bd9bfa960/coverage-7.13.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d490ba50c3f35dd7c17953c68f3270e7ccd1c6642e2d2afe2d8e720b98f5a053", size = 219278, upload-time = "2026-02-09T12:56:31.673Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/17/0cb7ca3de72e5f4ef2ec2fa0089beafbcaaaead1844e8b8a63d35173d77d/coverage-7.13.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:19bc3c88078789f8ef36acb014d7241961dbf883fd2533d18cb1e7a5b4e28b11", size = 219783, upload-time = "2026-02-09T12:56:33.104Z" },
+ { url = "https://files.pythonhosted.org/packages/ab/63/325d8e5b11e0eaf6d0f6a44fad444ae58820929a9b0de943fa377fe73e85/coverage-7.13.4-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:3998e5a32e62fdf410c0dbd3115df86297995d6e3429af80b8798aad894ca7aa", size = 250200, upload-time = "2026-02-09T12:56:34.474Z" },
+ { url = "https://files.pythonhosted.org/packages/76/53/c16972708cbb79f2942922571a687c52bd109a7bd51175aeb7558dff2236/coverage-7.13.4-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8e264226ec98e01a8e1054314af91ee6cde0eacac4f465cc93b03dbe0bce2fd7", size = 252114, upload-time = "2026-02-09T12:56:35.749Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/c2/7ab36d8b8cc412bec9ea2d07c83c48930eb4ba649634ba00cb7e4e0f9017/coverage-7.13.4-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a3aa4e7b9e416774b21797365b358a6e827ffadaaca81b69ee02946852449f00", size = 254220, upload-time = "2026-02-09T12:56:37.796Z" },
+ { url = "https://files.pythonhosted.org/packages/d6/4d/cf52c9a3322c89a0e6febdfbc83bb45c0ed3c64ad14081b9503adee702e7/coverage-7.13.4-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:71ca20079dd8f27fcf808817e281e90220475cd75115162218d0e27549f95fef", size = 256164, upload-time = "2026-02-09T12:56:39.016Z" },
+ { url = "https://files.pythonhosted.org/packages/78/e9/eb1dd17bd6de8289df3580e967e78294f352a5df8a57ff4671ee5fc3dcd0/coverage-7.13.4-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e2f25215f1a359ab17320b47bcdaca3e6e6356652e8256f2441e4ef972052903", size = 250325, upload-time = "2026-02-09T12:56:40.668Z" },
+ { url = "https://files.pythonhosted.org/packages/71/07/8c1542aa873728f72267c07278c5cc0ec91356daf974df21335ccdb46368/coverage-7.13.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d65b2d373032411e86960604dc4edac91fdfb5dca539461cf2cbe78327d1e64f", size = 251913, upload-time = "2026-02-09T12:56:41.97Z" },
+ { url = "https://files.pythonhosted.org/packages/74/d7/c62e2c5e4483a748e27868e4c32ad3daa9bdddbba58e1bc7a15e252baa74/coverage-7.13.4-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:94eb63f9b363180aff17de3e7c8760c3ba94664ea2695c52f10111244d16a299", size = 249974, upload-time = "2026-02-09T12:56:43.323Z" },
+ { url = "https://files.pythonhosted.org/packages/98/9f/4c5c015a6e98ced54efd0f5cf8d31b88e5504ecb6857585fc0161bb1e600/coverage-7.13.4-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:e856bf6616714c3a9fbc270ab54103f4e685ba236fa98c054e8f87f266c93505", size = 253741, upload-time = "2026-02-09T12:56:45.155Z" },
+ { url = "https://files.pythonhosted.org/packages/bd/59/0f4eef89b9f0fcd9633b5d350016f54126ab49426a70ff4c4e87446cabdc/coverage-7.13.4-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:65dfcbe305c3dfe658492df2d85259e0d79ead4177f9ae724b6fb245198f55d6", size = 249695, upload-time = "2026-02-09T12:56:46.636Z" },
+ { url = "https://files.pythonhosted.org/packages/b5/2c/b7476f938deb07166f3eb281a385c262675d688ff4659ad56c6c6b8e2e70/coverage-7.13.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b507778ae8a4c915436ed5c2e05b4a6cecfa70f734e19c22a005152a11c7b6a9", size = 250599, upload-time = "2026-02-09T12:56:48.13Z" },
+ { url = "https://files.pythonhosted.org/packages/b8/34/c3420709d9846ee3785b9f2831b4d94f276f38884032dca1457fa83f7476/coverage-7.13.4-cp311-cp311-win32.whl", hash = "sha256:784fc3cf8be001197b652d51d3fd259b1e2262888693a4636e18879f613a62a9", size = 221780, upload-time = "2026-02-09T12:56:50.479Z" },
+ { url = "https://files.pythonhosted.org/packages/61/08/3d9c8613079d2b11c185b865de9a4c1a68850cfda2b357fae365cf609f29/coverage-7.13.4-cp311-cp311-win_amd64.whl", hash = "sha256:2421d591f8ca05b308cf0092807308b2facbefe54af7c02ac22548b88b95c98f", size = 222715, upload-time = "2026-02-09T12:56:51.815Z" },
+ { url = "https://files.pythonhosted.org/packages/18/1a/54c3c80b2f056164cc0a6cdcb040733760c7c4be9d780fe655f356f433e4/coverage-7.13.4-cp311-cp311-win_arm64.whl", hash = "sha256:79e73a76b854d9c6088fe5d8b2ebe745f8681c55f7397c3c0a016192d681045f", size = 221385, upload-time = "2026-02-09T12:56:53.194Z" },
+ { url = "https://files.pythonhosted.org/packages/d1/81/4ce2fdd909c5a0ed1f6dedb88aa57ab79b6d1fbd9b588c1ac7ef45659566/coverage-7.13.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:02231499b08dabbe2b96612993e5fc34217cdae907a51b906ac7fca8027a4459", size = 219449, upload-time = "2026-02-09T12:56:54.889Z" },
+ { url = "https://files.pythonhosted.org/packages/5d/96/5238b1efc5922ddbdc9b0db9243152c09777804fb7c02ad1741eb18a11c0/coverage-7.13.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40aa8808140e55dc022b15d8aa7f651b6b3d68b365ea0398f1441e0b04d859c3", size = 219810, upload-time = "2026-02-09T12:56:56.33Z" },
+ { url = "https://files.pythonhosted.org/packages/78/72/2f372b726d433c9c35e56377cf1d513b4c16fe51841060d826b95caacec1/coverage-7.13.4-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5b856a8ccf749480024ff3bd7310adaef57bf31fd17e1bfc404b7940b6986634", size = 251308, upload-time = "2026-02-09T12:56:57.858Z" },
+ { url = "https://files.pythonhosted.org/packages/5d/a0/2ea570925524ef4e00bb6c82649f5682a77fac5ab910a65c9284de422600/coverage-7.13.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2c048ea43875fbf8b45d476ad79f179809c590ec7b79e2035c662e7afa3192e3", size = 254052, upload-time = "2026-02-09T12:56:59.754Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/ac/45dc2e19a1939098d783c846e130b8f862fbb50d09e0af663988f2f21973/coverage-7.13.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b7b38448866e83176e28086674fe7368ab8590e4610fb662b44e345b86d63ffa", size = 255165, upload-time = "2026-02-09T12:57:01.287Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/4d/26d236ff35abc3b5e63540d3386e4c3b192168c1d96da5cb2f43c640970f/coverage-7.13.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:de6defc1c9badbf8b9e67ae90fd00519186d6ab64e5cc5f3d21359c2a9b2c1d3", size = 257432, upload-time = "2026-02-09T12:57:02.637Z" },
+ { url = "https://files.pythonhosted.org/packages/ec/55/14a966c757d1348b2e19caf699415a2a4c4f7feaa4bbc6326a51f5c7dd1b/coverage-7.13.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7eda778067ad7ffccd23ecffce537dface96212576a07924cbf0d8799d2ded5a", size = 251716, upload-time = "2026-02-09T12:57:04.056Z" },
+ { url = "https://files.pythonhosted.org/packages/77/33/50116647905837c66d28b2af1321b845d5f5d19be9655cb84d4a0ea806b4/coverage-7.13.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e87f6c587c3f34356c3759f0420693e35e7eb0e2e41e4c011cb6ec6ecbbf1db7", size = 253089, upload-time = "2026-02-09T12:57:05.503Z" },
+ { url = "https://files.pythonhosted.org/packages/c2/b4/8efb11a46e3665d92635a56e4f2d4529de6d33f2cb38afd47d779d15fc99/coverage-7.13.4-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8248977c2e33aecb2ced42fef99f2d319e9904a36e55a8a68b69207fb7e43edc", size = 251232, upload-time = "2026-02-09T12:57:06.879Z" },
+ { url = "https://files.pythonhosted.org/packages/51/24/8cd73dd399b812cc76bb0ac260e671c4163093441847ffe058ac9fda1e32/coverage-7.13.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:25381386e80ae727608e662474db537d4df1ecd42379b5ba33c84633a2b36d47", size = 255299, upload-time = "2026-02-09T12:57:08.245Z" },
+ { url = "https://files.pythonhosted.org/packages/03/94/0a4b12f1d0e029ce1ccc1c800944a9984cbe7d678e470bb6d3c6bc38a0da/coverage-7.13.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:ee756f00726693e5ba94d6df2bdfd64d4852d23b09bb0bc700e3b30e6f333985", size = 250796, upload-time = "2026-02-09T12:57:10.142Z" },
+ { url = "https://files.pythonhosted.org/packages/73/44/6002fbf88f6698ca034360ce474c406be6d5a985b3fdb3401128031eef6b/coverage-7.13.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fdfc1e28e7c7cdce44985b3043bc13bbd9c747520f94a4d7164af8260b3d91f0", size = 252673, upload-time = "2026-02-09T12:57:12.197Z" },
+ { url = "https://files.pythonhosted.org/packages/de/c6/a0279f7c00e786be75a749a5674e6fa267bcbd8209cd10c9a450c655dfa7/coverage-7.13.4-cp312-cp312-win32.whl", hash = "sha256:01d4cbc3c283a17fc1e42d614a119f7f438eabb593391283adca8dc86eff1246", size = 221990, upload-time = "2026-02-09T12:57:14.085Z" },
+ { url = "https://files.pythonhosted.org/packages/77/4e/c0a25a425fcf5557d9abd18419c95b63922e897bc86c1f327f155ef234a9/coverage-7.13.4-cp312-cp312-win_amd64.whl", hash = "sha256:9401ebc7ef522f01d01d45532c68c5ac40fb27113019b6b7d8b208f6e9baa126", size = 222800, upload-time = "2026-02-09T12:57:15.944Z" },
+ { url = "https://files.pythonhosted.org/packages/47/ac/92da44ad9a6f4e3a7debd178949d6f3769bedca33830ce9b1dcdab589a37/coverage-7.13.4-cp312-cp312-win_arm64.whl", hash = "sha256:b1ec7b6b6e93255f952e27ab58fbc68dcc468844b16ecbee881aeb29b6ab4d8d", size = 221415, upload-time = "2026-02-09T12:57:17.497Z" },
+ { url = "https://files.pythonhosted.org/packages/db/23/aad45061a31677d68e47499197a131eea55da4875d16c1f42021ab963503/coverage-7.13.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b66a2da594b6068b48b2692f043f35d4d3693fb639d5ea8b39533c2ad9ac3ab9", size = 219474, upload-time = "2026-02-09T12:57:19.332Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/70/9b8b67a0945f3dfec1fd896c5cefb7c19d5a3a6d74630b99a895170999ae/coverage-7.13.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3599eb3992d814d23b35c536c28df1a882caa950f8f507cef23d1cbf334995ac", size = 219844, upload-time = "2026-02-09T12:57:20.66Z" },
+ { url = "https://files.pythonhosted.org/packages/97/fd/7e859f8fab324cef6c4ad7cff156ca7c489fef9179d5749b0c8d321281c2/coverage-7.13.4-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:93550784d9281e374fb5a12bf1324cc8a963fd63b2d2f223503ef0fd4aa339ea", size = 250832, upload-time = "2026-02-09T12:57:22.007Z" },
+ { url = "https://files.pythonhosted.org/packages/e4/dc/b2442d10020c2f52617828862d8b6ee337859cd8f3a1f13d607dddda9cf7/coverage-7.13.4-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b720ce6a88a2755f7c697c23268ddc47a571b88052e6b155224347389fdf6a3b", size = 253434, upload-time = "2026-02-09T12:57:23.339Z" },
+ { url = "https://files.pythonhosted.org/packages/5a/88/6728a7ad17428b18d836540630487231f5470fb82454871149502f5e5aa2/coverage-7.13.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7b322db1284a2ed3aa28ffd8ebe3db91c929b7a333c0820abec3d838ef5b3525", size = 254676, upload-time = "2026-02-09T12:57:24.774Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/bc/21244b1b8cedf0dff0a2b53b208015fe798d5f2a8d5348dbfece04224fff/coverage-7.13.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f4594c67d8a7c89cf922d9df0438c7c7bb022ad506eddb0fdb2863359ff78242", size = 256807, upload-time = "2026-02-09T12:57:26.125Z" },
+ { url = "https://files.pythonhosted.org/packages/97/a0/ddba7ed3251cff51006737a727d84e05b61517d1784a9988a846ba508877/coverage-7.13.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:53d133df809c743eb8bce33b24bcababb371f4441340578cd406e084d94a6148", size = 251058, upload-time = "2026-02-09T12:57:27.614Z" },
+ { url = "https://files.pythonhosted.org/packages/9b/55/e289addf7ff54d3a540526f33751951bf0878f3809b47f6dfb3def69c6f7/coverage-7.13.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:76451d1978b95ba6507a039090ba076105c87cc76fc3efd5d35d72093964d49a", size = 252805, upload-time = "2026-02-09T12:57:29.066Z" },
+ { url = "https://files.pythonhosted.org/packages/13/4e/cc276b1fa4a59be56d96f1dabddbdc30f4ba22e3b1cd42504c37b3313255/coverage-7.13.4-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:7f57b33491e281e962021de110b451ab8a24182589be17e12a22c79047935e23", size = 250766, upload-time = "2026-02-09T12:57:30.522Z" },
+ { url = "https://files.pythonhosted.org/packages/94/44/1093b8f93018f8b41a8cf29636c9292502f05e4a113d4d107d14a3acd044/coverage-7.13.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:1731dc33dc276dafc410a885cbf5992f1ff171393e48a21453b78727d090de80", size = 254923, upload-time = "2026-02-09T12:57:31.946Z" },
+ { url = "https://files.pythonhosted.org/packages/8b/55/ea2796da2d42257f37dbea1aab239ba9263b31bd91d5527cdd6db5efe174/coverage-7.13.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:bd60d4fe2f6fa7dff9223ca1bbc9f05d2b6697bc5961072e5d3b952d46e1b1ea", size = 250591, upload-time = "2026-02-09T12:57:33.842Z" },
+ { url = "https://files.pythonhosted.org/packages/d4/fa/7c4bb72aacf8af5020675aa633e59c1fbe296d22aed191b6a5b711eb2bc7/coverage-7.13.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9181a3ccead280b828fae232df12b16652702b49d41e99d657f46cc7b1f6ec7a", size = 252364, upload-time = "2026-02-09T12:57:35.743Z" },
+ { url = "https://files.pythonhosted.org/packages/5c/38/a8d2ec0146479c20bbaa7181b5b455a0c41101eed57f10dd19a78ab44c80/coverage-7.13.4-cp313-cp313-win32.whl", hash = "sha256:f53d492307962561ac7de4cd1de3e363589b000ab69617c6156a16ba7237998d", size = 222010, upload-time = "2026-02-09T12:57:37.25Z" },
+ { url = "https://files.pythonhosted.org/packages/e2/0c/dbfafbe90a185943dcfbc766fe0e1909f658811492d79b741523a414a6cc/coverage-7.13.4-cp313-cp313-win_amd64.whl", hash = "sha256:e6f70dec1cc557e52df5306d051ef56003f74d56e9c4dd7ddb07e07ef32a84dd", size = 222818, upload-time = "2026-02-09T12:57:38.734Z" },
+ { url = "https://files.pythonhosted.org/packages/04/d1/934918a138c932c90d78301f45f677fb05c39a3112b96fd2c8e60503cdc7/coverage-7.13.4-cp313-cp313-win_arm64.whl", hash = "sha256:fb07dc5da7e849e2ad31a5d74e9bece81f30ecf5a42909d0a695f8bd1874d6af", size = 221438, upload-time = "2026-02-09T12:57:40.223Z" },
+ { url = "https://files.pythonhosted.org/packages/52/57/ee93ced533bcb3e6df961c0c6e42da2fc6addae53fb95b94a89b1e33ebd7/coverage-7.13.4-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:40d74da8e6c4b9ac18b15331c4b5ebc35a17069410cad462ad4f40dcd2d50c0d", size = 220165, upload-time = "2026-02-09T12:57:41.639Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/e0/969fc285a6fbdda49d91af278488d904dcd7651b2693872f0ff94e40e84a/coverage-7.13.4-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4223b4230a376138939a9173f1bdd6521994f2aff8047fae100d6d94d50c5a12", size = 220516, upload-time = "2026-02-09T12:57:44.215Z" },
+ { url = "https://files.pythonhosted.org/packages/b1/b8/9531944e16267e2735a30a9641ff49671f07e8138ecf1ca13db9fd2560c7/coverage-7.13.4-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:1d4be36a5114c499f9f1f9195e95ebf979460dbe2d88e6816ea202010ba1c34b", size = 261804, upload-time = "2026-02-09T12:57:45.989Z" },
+ { url = "https://files.pythonhosted.org/packages/8a/f3/e63df6d500314a2a60390d1989240d5f27318a7a68fa30ad3806e2a9323e/coverage-7.13.4-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:200dea7d1e8095cc6e98cdabe3fd1d21ab17d3cee6dab00cadbb2fe35d9c15b9", size = 263885, upload-time = "2026-02-09T12:57:47.42Z" },
+ { url = "https://files.pythonhosted.org/packages/f3/67/7654810de580e14b37670b60a09c599fa348e48312db5b216d730857ffe6/coverage-7.13.4-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b8eb931ee8e6d8243e253e5ed7336deea6904369d2fd8ae6e43f68abbf167092", size = 266308, upload-time = "2026-02-09T12:57:49.345Z" },
+ { url = "https://files.pythonhosted.org/packages/37/6f/39d41eca0eab3cc82115953ad41c4e77935286c930e8fad15eaed1389d83/coverage-7.13.4-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:75eab1ebe4f2f64d9509b984f9314d4aa788540368218b858dad56dc8f3e5eb9", size = 267452, upload-time = "2026-02-09T12:57:50.811Z" },
+ { url = "https://files.pythonhosted.org/packages/50/6d/39c0fbb8fc5cd4d2090811e553c2108cf5112e882f82505ee7495349a6bf/coverage-7.13.4-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c35eb28c1d085eb7d8c9b3296567a1bebe03ce72962e932431b9a61f28facf26", size = 261057, upload-time = "2026-02-09T12:57:52.447Z" },
+ { url = "https://files.pythonhosted.org/packages/a4/a2/60010c669df5fa603bb5a97fb75407e191a846510da70ac657eb696b7fce/coverage-7.13.4-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:eb88b316ec33760714a4720feb2816a3a59180fd58c1985012054fa7aebee4c2", size = 263875, upload-time = "2026-02-09T12:57:53.938Z" },
+ { url = "https://files.pythonhosted.org/packages/3e/d9/63b22a6bdbd17f1f96e9ed58604c2a6b0e72a9133e37d663bef185877cf6/coverage-7.13.4-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:7d41eead3cc673cbd38a4417deb7fd0b4ca26954ff7dc6078e33f6ff97bed940", size = 261500, upload-time = "2026-02-09T12:57:56.012Z" },
+ { url = "https://files.pythonhosted.org/packages/70/bf/69f86ba1ad85bc3ad240e4c0e57a2e620fbc0e1645a47b5c62f0e941ad7f/coverage-7.13.4-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:fb26a934946a6afe0e326aebe0730cdff393a8bc0bbb65a2f41e30feddca399c", size = 265212, upload-time = "2026-02-09T12:57:57.5Z" },
+ { url = "https://files.pythonhosted.org/packages/ae/f2/5f65a278a8c2148731831574c73e42f57204243d33bedaaf18fa79c5958f/coverage-7.13.4-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:dae88bc0fc77edaa65c14be099bd57ee140cf507e6bfdeea7938457ab387efb0", size = 260398, upload-time = "2026-02-09T12:57:59.027Z" },
+ { url = "https://files.pythonhosted.org/packages/ef/80/6e8280a350ee9fea92f14b8357448a242dcaa243cb2c72ab0ca591f66c8c/coverage-7.13.4-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:845f352911777a8e722bfce168958214951e07e47e5d5d9744109fa5fe77f79b", size = 262584, upload-time = "2026-02-09T12:58:01.129Z" },
+ { url = "https://files.pythonhosted.org/packages/22/63/01ff182fc95f260b539590fb12c11ad3e21332c15f9799cb5e2386f71d9f/coverage-7.13.4-cp313-cp313t-win32.whl", hash = "sha256:2fa8d5f8de70688a28240de9e139fa16b153cc3cbb01c5f16d88d6505ebdadf9", size = 222688, upload-time = "2026-02-09T12:58:02.736Z" },
+ { url = "https://files.pythonhosted.org/packages/a9/43/89de4ef5d3cd53b886afa114065f7e9d3707bdb3e5efae13535b46ae483d/coverage-7.13.4-cp313-cp313t-win_amd64.whl", hash = "sha256:9351229c8c8407645840edcc277f4a2d44814d1bc34a2128c11c2a031d45a5dd", size = 223746, upload-time = "2026-02-09T12:58:05.362Z" },
+ { url = "https://files.pythonhosted.org/packages/35/39/7cf0aa9a10d470a5309b38b289b9bb07ddeac5d61af9b664fe9775a4cb3e/coverage-7.13.4-cp313-cp313t-win_arm64.whl", hash = "sha256:30b8d0512f2dc8c8747557e8fb459d6176a2c9e5731e2b74d311c03b78451997", size = 222003, upload-time = "2026-02-09T12:58:06.952Z" },
+ { url = "https://files.pythonhosted.org/packages/92/11/a9cf762bb83386467737d32187756a42094927150c3e107df4cb078e8590/coverage-7.13.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:300deaee342f90696ed186e3a00c71b5b3d27bffe9e827677954f4ee56969601", size = 219522, upload-time = "2026-02-09T12:58:08.623Z" },
+ { url = "https://files.pythonhosted.org/packages/d3/28/56e6d892b7b052236d67c95f1936b6a7cf7c3e2634bf27610b8cbd7f9c60/coverage-7.13.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:29e3220258d682b6226a9b0925bc563ed9a1ebcff3cad30f043eceea7eaf2689", size = 219855, upload-time = "2026-02-09T12:58:10.176Z" },
+ { url = "https://files.pythonhosted.org/packages/e5/69/233459ee9eb0c0d10fcc2fe425a029b3fa5ce0f040c966ebce851d030c70/coverage-7.13.4-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:391ee8f19bef69210978363ca930f7328081c6a0152f1166c91f0b5fdd2a773c", size = 250887, upload-time = "2026-02-09T12:58:12.503Z" },
+ { url = "https://files.pythonhosted.org/packages/06/90/2cdab0974b9b5bbc1623f7876b73603aecac11b8d95b85b5b86b32de5eab/coverage-7.13.4-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:0dd7ab8278f0d58a0128ba2fca25824321f05d059c1441800e934ff2efa52129", size = 253396, upload-time = "2026-02-09T12:58:14.615Z" },
+ { url = "https://files.pythonhosted.org/packages/ac/15/ea4da0f85bf7d7b27635039e649e99deb8173fe551096ea15017f7053537/coverage-7.13.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78cdf0d578b15148b009ccf18c686aa4f719d887e76e6b40c38ffb61d264a552", size = 254745, upload-time = "2026-02-09T12:58:16.162Z" },
+ { url = "https://files.pythonhosted.org/packages/99/11/bb356e86920c655ca4d61daee4e2bbc7258f0a37de0be32d233b561134ff/coverage-7.13.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:48685fee12c2eb3b27c62f2658e7ea21e9c3239cba5a8a242801a0a3f6a8c62a", size = 257055, upload-time = "2026-02-09T12:58:17.892Z" },
+ { url = "https://files.pythonhosted.org/packages/c9/0f/9ae1f8cb17029e09da06ca4e28c9e1d5c1c0a511c7074592e37e0836c915/coverage-7.13.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4e83efc079eb39480e6346a15a1bcb3e9b04759c5202d157e1dd4303cd619356", size = 250911, upload-time = "2026-02-09T12:58:19.495Z" },
+ { url = "https://files.pythonhosted.org/packages/89/3a/adfb68558fa815cbc29747b553bc833d2150228f251b127f1ce97e48547c/coverage-7.13.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ecae9737b72408d6a950f7e525f30aca12d4bd8dd95e37342e5beb3a2a8c4f71", size = 252754, upload-time = "2026-02-09T12:58:21.064Z" },
+ { url = "https://files.pythonhosted.org/packages/32/b1/540d0c27c4e748bd3cd0bd001076ee416eda993c2bae47a73b7cc9357931/coverage-7.13.4-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:ae4578f8528569d3cf303fef2ea569c7f4c4059a38c8667ccef15c6e1f118aa5", size = 250720, upload-time = "2026-02-09T12:58:22.622Z" },
+ { url = "https://files.pythonhosted.org/packages/c7/95/383609462b3ffb1fe133014a7c84fc0dd01ed55ac6140fa1093b5af7ebb1/coverage-7.13.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:6fdef321fdfbb30a197efa02d48fcd9981f0d8ad2ae8903ac318adc653f5df98", size = 254994, upload-time = "2026-02-09T12:58:24.548Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/ba/1761138e86c81680bfc3c49579d66312865457f9fe405b033184e5793cb3/coverage-7.13.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:2b0f6ccf3dbe577170bebfce1318707d0e8c3650003cb4b3a9dd744575daa8b5", size = 250531, upload-time = "2026-02-09T12:58:26.271Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/8e/05900df797a9c11837ab59c4d6fe94094e029582aab75c3309a93e6fb4e3/coverage-7.13.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:75fcd519f2a5765db3f0e391eb3b7d150cce1a771bf4c9f861aeab86c767a3c0", size = 252189, upload-time = "2026-02-09T12:58:27.807Z" },
+ { url = "https://files.pythonhosted.org/packages/00/bd/29c9f2db9ea4ed2738b8a9508c35626eb205d51af4ab7bf56a21a2e49926/coverage-7.13.4-cp314-cp314-win32.whl", hash = "sha256:8e798c266c378da2bd819b0677df41ab46d78065fb2a399558f3f6cae78b2fbb", size = 222258, upload-time = "2026-02-09T12:58:29.441Z" },
+ { url = "https://files.pythonhosted.org/packages/a7/4d/1f8e723f6829977410efeb88f73673d794075091c8c7c18848d273dc9d73/coverage-7.13.4-cp314-cp314-win_amd64.whl", hash = "sha256:245e37f664d89861cf2329c9afa2c1fe9e6d4e1a09d872c947e70718aeeac505", size = 223073, upload-time = "2026-02-09T12:58:31.026Z" },
+ { url = "https://files.pythonhosted.org/packages/51/5b/84100025be913b44e082ea32abcf1afbf4e872f5120b7a1cab1d331b1e13/coverage-7.13.4-cp314-cp314-win_arm64.whl", hash = "sha256:ad27098a189e5838900ce4c2a99f2fe42a0bf0c2093c17c69b45a71579e8d4a2", size = 221638, upload-time = "2026-02-09T12:58:32.599Z" },
+ { url = "https://files.pythonhosted.org/packages/a7/e4/c884a405d6ead1370433dad1e3720216b4f9fd8ef5b64bfd984a2a60a11a/coverage-7.13.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:85480adfb35ffc32d40918aad81b89c69c9cc5661a9b8a81476d3e645321a056", size = 220246, upload-time = "2026-02-09T12:58:34.181Z" },
+ { url = "https://files.pythonhosted.org/packages/81/5c/4d7ed8b23b233b0fffbc9dfec53c232be2e695468523242ea9fd30f97ad2/coverage-7.13.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:79be69cf7f3bf9b0deeeb062eab7ac7f36cd4cc4c4dd694bd28921ba4d8596cc", size = 220514, upload-time = "2026-02-09T12:58:35.704Z" },
+ { url = "https://files.pythonhosted.org/packages/2f/6f/3284d4203fd2f28edd73034968398cd2d4cb04ab192abc8cff007ea35679/coverage-7.13.4-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:caa421e2684e382c5d8973ac55e4f36bed6821a9bad5c953494de960c74595c9", size = 261877, upload-time = "2026-02-09T12:58:37.864Z" },
+ { url = "https://files.pythonhosted.org/packages/09/aa/b672a647bbe1556a85337dc95bfd40d146e9965ead9cc2fe81bde1e5cbce/coverage-7.13.4-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:14375934243ee05f56c45393fe2ce81fe5cc503c07cee2bdf1725fb8bef3ffaf", size = 264004, upload-time = "2026-02-09T12:58:39.492Z" },
+ { url = "https://files.pythonhosted.org/packages/79/a1/aa384dbe9181f98bba87dd23dda436f0c6cf2e148aecbb4e50fc51c1a656/coverage-7.13.4-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:25a41c3104d08edb094d9db0d905ca54d0cd41c928bb6be3c4c799a54753af55", size = 266408, upload-time = "2026-02-09T12:58:41.852Z" },
+ { url = "https://files.pythonhosted.org/packages/53/5e/5150bf17b4019bc600799f376bb9606941e55bd5a775dc1e096b6ffea952/coverage-7.13.4-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6f01afcff62bf9a08fb32b2c1d6e924236c0383c02c790732b6537269e466a72", size = 267544, upload-time = "2026-02-09T12:58:44.093Z" },
+ { url = "https://files.pythonhosted.org/packages/e0/ed/f1de5c675987a4a7a672250d2c5c9d73d289dbf13410f00ed7181d8017dd/coverage-7.13.4-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:eb9078108fbf0bcdde37c3f4779303673c2fa1fe8f7956e68d447d0dd426d38a", size = 260980, upload-time = "2026-02-09T12:58:45.721Z" },
+ { url = "https://files.pythonhosted.org/packages/b3/e3/fe758d01850aa172419a6743fe76ba8b92c29d181d4f676ffe2dae2ba631/coverage-7.13.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0e086334e8537ddd17e5f16a344777c1ab8194986ec533711cbe6c41cde841b6", size = 263871, upload-time = "2026-02-09T12:58:47.334Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/76/b829869d464115e22499541def9796b25312b8cf235d3bb00b39f1675395/coverage-7.13.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:725d985c5ab621268b2edb8e50dfe57633dc69bda071abc470fed55a14935fd3", size = 261472, upload-time = "2026-02-09T12:58:48.995Z" },
+ { url = "https://files.pythonhosted.org/packages/14/9e/caedb1679e73e2f6ad240173f55218488bfe043e38da577c4ec977489915/coverage-7.13.4-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:3c06f0f1337c667b971ca2f975523347e63ec5e500b9aa5882d91931cd3ef750", size = 265210, upload-time = "2026-02-09T12:58:51.178Z" },
+ { url = "https://files.pythonhosted.org/packages/3a/10/0dd02cb009b16ede425b49ec344aba13a6ae1dc39600840ea6abcb085ac4/coverage-7.13.4-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:590c0ed4bf8e85f745e6b805b2e1c457b2e33d5255dd9729743165253bc9ad39", size = 260319, upload-time = "2026-02-09T12:58:53.081Z" },
+ { url = "https://files.pythonhosted.org/packages/92/8e/234d2c927af27c6d7a5ffad5bd2cf31634c46a477b4c7adfbfa66baf7ebb/coverage-7.13.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:eb30bf180de3f632cd043322dad5751390e5385108b2807368997d1a92a509d0", size = 262638, upload-time = "2026-02-09T12:58:55.258Z" },
+ { url = "https://files.pythonhosted.org/packages/2f/64/e5547c8ff6964e5965c35a480855911b61509cce544f4d442caa759a0702/coverage-7.13.4-cp314-cp314t-win32.whl", hash = "sha256:c4240e7eded42d131a2d2c4dec70374b781b043ddc79a9de4d55ca71f8e98aea", size = 223040, upload-time = "2026-02-09T12:58:56.936Z" },
+ { url = "https://files.pythonhosted.org/packages/c7/96/38086d58a181aac86d503dfa9c47eb20715a79c3e3acbdf786e92e5c09a8/coverage-7.13.4-cp314-cp314t-win_amd64.whl", hash = "sha256:4c7d3cc01e7350f2f0f6f7036caaf5673fb56b6998889ccfe9e1c1fe75a9c932", size = 224148, upload-time = "2026-02-09T12:58:58.645Z" },
+ { url = "https://files.pythonhosted.org/packages/ce/72/8d10abd3740a0beb98c305e0c3faf454366221c0f37a8bcf8f60020bb65a/coverage-7.13.4-cp314-cp314t-win_arm64.whl", hash = "sha256:23e3f687cf945070d1c90f85db66d11e3025665d8dafa831301a0e0038f3db9b", size = 222172, upload-time = "2026-02-09T12:59:00.396Z" },
+ { url = "https://files.pythonhosted.org/packages/0d/4a/331fe2caf6799d591109bb9c08083080f6de90a823695d412a935622abb2/coverage-7.13.4-py3-none-any.whl", hash = "sha256:1af1641e57cf7ba1bd67d677c9abdbcd6cc2ab7da3bca7fa1e2b7e50e65f2ad0", size = 211242, upload-time = "2026-02-09T12:59:02.032Z" },
+]
+
+[package.optional-dependencies]
+toml = [
+ { name = "tomli", marker = "python_full_version <= '3.11'" },
+]
+
[[package]]
name = "croniter"
version = "6.0.0"
@@ -1767,6 +1885,7 @@ dev = [
{ name = "pre-commit" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
+ { name = "pytest-cov" },
{ name = "ruff" },
]
full = [
@@ -1806,6 +1925,7 @@ requires-dist = [
{ name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.0.0,<3.0.0" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0,<10.0.0" },
{ name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0,<2.0.0" },
+ { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=6.0.0" },
{ name = "python-dotenv", specifier = ">=0.1.0,<2.0.0" },
{ name = "pyyaml", specifier = ">=6.0,<7.0" },
{ name = "requests", specifier = ">=2.0.0,<3.0.0" },
@@ -2349,6 +2469,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
]
+[[package]]
+name = "pytest-cov"
+version = "6.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "coverage", extra = ["toml"] },
+ { name = "pluggy" },
+ { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/30/4c/f883ab8f0daad69f47efdf95f55a66b51a8b939c430dadce0611508d9e99/pytest_cov-6.3.0.tar.gz", hash = "sha256:35c580e7800f87ce892e687461166e1ac2bcb8fb9e13aea79032518d6e503ff2", size = 70398, upload-time = "2025-09-06T15:40:14.361Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/80/b4/bb7263e12aade3842b938bc5c6958cae79c5ee18992f9b9349019579da0f/pytest_cov-6.3.0-py3-none-any.whl", hash = "sha256:440db28156d2468cafc0415b4f8e50856a0d11faefa38f30906048fe490f1749", size = 25115, upload-time = "2025-09-06T15:40:12.44Z" },
+]
+
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"
From 3110852ead43062a5b531e84ac4c4b9543ed1d0b Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 14:18:46 -0800
Subject: [PATCH 14/60] Enhance test suite with new agent loop stubs and
utility functions
- Introduced `SimpleAgentLoop`, `FailingAgentLoop`, and `SlowAgentLoop` classes for testing purposes in `conftest.py`.
- Added `make_rollout_payload` and `mock_llm_client` utility functions to streamline test setup.
- Updated various test files to utilize the new agent loop stubs and utility functions, improving test coverage and maintainability.
- Refactored existing tests to use parameterized inputs for better clarity and efficiency.
---
tests/conftest.py | 84 ++++
tests/test_auth_flow.py | 495 +++++++++++++++++++++++
tests/test_eval_runner.py | 665 +++++++++++++++++++++++++++++++
tests/test_rollout_base.py | 9 +-
tests/test_rollout_exceptions.py | 40 +-
tests/test_rollout_schemas.py | 107 +----
tests/test_rollout_server.py | 47 +--
tests/test_rollout_server_app.py | 197 ++-------
tests/test_rollout_state.py | 599 ++++++++++++++++++++++++++++
tests/test_rollout_utils.py | 134 +++----
tests/test_settings.py | 576 ++++++++++++++++++++++++++
11 files changed, 2540 insertions(+), 413 deletions(-)
create mode 100644 tests/test_auth_flow.py
create mode 100644 tests/test_rollout_state.py
create mode 100644 tests/test_settings.py
diff --git a/tests/conftest.py b/tests/conftest.py
index e64219af..10699cac 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,7 +2,9 @@
from __future__ import annotations
+import asyncio
from typing import Any
+from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -13,6 +15,7 @@
OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
+ RolloutMetrics,
RolloutRequest,
RolloutResult,
)
@@ -123,3 +126,84 @@ def sample_tool_message() -> dict[str, Any]:
"content": "4",
"tool_call_id": "call_123",
}
+
+
+# =============================================================================
+# Shared agent loop stubs (used by multiple test modules)
+# =============================================================================
+
+
+class SimpleAgentLoop(RolloutAgentLoop):
+ """Simple agent loop that completes immediately."""
+
+ name = "simple_agent"
+
+ def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
+ self._tools = tools or []
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return self._tools
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ return ctx.complete(list(ctx.request.messages))
+
+
+class FailingAgentLoop(RolloutAgentLoop):
+ """Agent loop that raises an error."""
+
+ name = "failing_agent"
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ raise RuntimeError("Test error")
+
+
+class SlowAgentLoop(RolloutAgentLoop):
+ """Agent loop that sleeps for a configurable duration before completing."""
+
+ name = "slow_agent"
+
+ def __init__(self, delay: float = 0.5):
+ self._delay = delay
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ await asyncio.sleep(self._delay)
+ return ctx.complete(list(ctx.request.messages))
+
+
+# =============================================================================
+# Shared helper functions (used by multiple test modules)
+# =============================================================================
+
+
+def make_rollout_payload(
+ rollout_id: str = "test-123",
+ idempotency_key: str | None = None,
+ **overrides: Any,
+) -> dict[str, Any]:
+ """Build a minimal valid RolloutRequest JSON payload."""
+ payload: dict[str, Any] = {
+ "rollout_id": rollout_id,
+ "server_url": "http://localhost:8080",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "completion_params": {"temperature": 0.7},
+ }
+ if idempotency_key is not None:
+ payload["idempotency_key"] = idempotency_key
+ payload.update(overrides)
+ return payload
+
+
+def mock_llm_client() -> MagicMock:
+ """Create a mocked OsmosisLLMClient usable as an async context manager."""
+ mock = AsyncMock()
+ mock.__aenter__ = AsyncMock(return_value=mock)
+ mock.__aexit__ = AsyncMock(return_value=None)
+ mock.complete_rollout = AsyncMock()
+ mock.get_metrics = MagicMock(return_value=RolloutMetrics())
+ return mock
diff --git a/tests/test_auth_flow.py b/tests/test_auth_flow.py
new file mode 100644
index 00000000..b4e600b0
--- /dev/null
+++ b/tests/test_auth_flow.py
@@ -0,0 +1,495 @@
+"""Tests for osmosis_ai.auth.flow - login orchestration."""
+
+from __future__ import annotations
+
+import json
+import socket
+from datetime import datetime, timedelta, timezone
+from typing import Any
+from unittest.mock import MagicMock, patch
+from urllib.error import HTTPError, URLError
+from urllib.parse import parse_qs, urlparse
+
+import pytest
+
+from osmosis_ai.auth.config import PLATFORM_URL
+from osmosis_ai.auth.credentials import OrganizationInfo, UserInfo, WorkspaceCredentials
+from osmosis_ai.auth.flow import (
+ LoginError,
+ LoginResult,
+ _build_login_url,
+ _generate_state,
+ _get_device_name,
+ _verify_and_get_user_info,
+ login,
+)
+
+# ---------------------------------------------------------------------------
+# _generate_state
+# ---------------------------------------------------------------------------
+
+
+class TestGenerateState:
+ def test_returns_string(self) -> None:
+ result = _generate_state()
+ assert isinstance(result, str)
+
+ def test_returns_non_empty(self) -> None:
+ result = _generate_state()
+ assert len(result) > 0
+
+ def test_returns_unique_values(self) -> None:
+ states = {_generate_state() for _ in range(50)}
+ assert len(states) == 50, "Each generated state should be unique"
+
+
+# ---------------------------------------------------------------------------
+# _get_device_name
+# ---------------------------------------------------------------------------
+
+
+class TestGetDeviceName:
+ def test_returns_hostname(self) -> None:
+ name = _get_device_name()
+ assert isinstance(name, str)
+ assert len(name) > 0
+
+ def test_matches_socket_gethostname(self) -> None:
+ expected = socket.gethostname()
+ assert _get_device_name() == expected
+
+ def test_falls_back_on_exception(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(
+ socket, "gethostname", lambda: (_ for _ in ()).throw(OSError("mocked"))
+ )
+ assert _get_device_name() == "Unknown"
+
+
+# ---------------------------------------------------------------------------
+# _build_login_url
+# ---------------------------------------------------------------------------
+
+
+class TestBuildLoginUrl:
+ def test_url_starts_with_platform_url(self) -> None:
+ url = _build_login_url(state="abc123", port=8976)
+ assert url.startswith(PLATFORM_URL)
+
+ def test_url_path_is_cli_auth(self) -> None:
+ url = _build_login_url(state="abc123", port=8976)
+ parsed = urlparse(url)
+ assert parsed.path == "/cli-auth"
+
+ def test_query_params_include_state(self) -> None:
+ url = _build_login_url(state="mystate", port=8976)
+ parsed = urlparse(url)
+ params = parse_qs(parsed.query)
+ assert params["state"] == ["mystate"]
+
+ def test_query_params_include_port(self) -> None:
+ url = _build_login_url(state="s", port=9999)
+ parsed = urlparse(url)
+ params = parse_qs(parsed.query)
+ assert params["port"] == ["9999"]
+
+ def test_query_params_include_redirect_uri(self) -> None:
+ url = _build_login_url(state="s", port=8980)
+ parsed = urlparse(url)
+ params = parse_qs(parsed.query)
+ assert params["redirect_uri"] == ["http://localhost:8980/callback"]
+
+ def test_query_params_include_device_name(self) -> None:
+ url = _build_login_url(state="s", port=8976)
+ parsed = urlparse(url)
+ params = parse_qs(parsed.query)
+ assert "device_name" in params
+ assert len(params["device_name"][0]) > 0
+
+
+# ---------------------------------------------------------------------------
+# _verify_and_get_user_info
+# ---------------------------------------------------------------------------
+
+
+def _make_verify_response(
+ *,
+ valid: bool = True,
+ user: dict[str, Any] | None = None,
+ organization: dict[str, Any] | None = None,
+ expires_at: str | None = None,
+ token_id: str | None = "tok_abc",
+) -> bytes:
+ """Build a JSON response body for the /api/cli/verify endpoint."""
+ data: dict[str, Any] = {"valid": valid}
+ if user is not None:
+ data["user"] = user
+ else:
+ data["user"] = {"id": "u1", "email": "u@test.com", "name": "Test User"}
+ if organization is not None:
+ data["organization"] = organization
+ else:
+ data["organization"] = {"id": "o1", "name": "TestOrg", "role": "owner"}
+ if expires_at is not None:
+ data["expires_at"] = expires_at
+ if token_id is not None:
+ data["token_id"] = token_id
+ return json.dumps(data).encode()
+
+
+class TestVerifyAndGetUserInfo:
+ def test_successful_verification(self) -> None:
+ expires_str = (datetime.now(timezone.utc) + timedelta(days=90)).isoformat()
+ body = _make_verify_response(expires_at=expires_str)
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ user, org, exp, tid = _verify_and_get_user_info("test-token")
+
+ assert isinstance(user, UserInfo)
+ assert user.email == "u@test.com"
+ assert isinstance(org, OrganizationInfo)
+ assert org.name == "TestOrg"
+ assert exp.tzinfo is not None
+ assert tid == "tok_abc"
+
+ def test_verification_with_no_expires_at_defaults_to_90_days(self) -> None:
+ body = _make_verify_response(expires_at=None)
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ before = datetime.now(timezone.utc)
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ _, _, exp, _ = _verify_and_get_user_info("token")
+ after = datetime.now(timezone.utc)
+
+ assert exp >= before + timedelta(days=89)
+ assert exp <= after + timedelta(days=91)
+
+ def test_invalid_token_raises_login_error(self) -> None:
+ body = _make_verify_response(valid=False)
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ with pytest.raises(LoginError, match="Token verification failed"):
+ _verify_and_get_user_info("bad-token")
+
+ def test_http_401_raises_login_error(self) -> None:
+ error = HTTPError(
+ url="http://test",
+ code=401,
+ msg="Unauthorized",
+ hdrs=None,
+ fp=None, # type: ignore[arg-type]
+ )
+ with patch("urllib.request.urlopen", side_effect=error):
+ with pytest.raises(LoginError, match="Invalid or expired token"):
+ _verify_and_get_user_info("expired-token")
+
+ def test_http_500_raises_login_error(self) -> None:
+ error = HTTPError(
+ url="http://test",
+ code=500,
+ msg="Server Error",
+ hdrs=None,
+ fp=None, # type: ignore[arg-type]
+ )
+ with patch("urllib.request.urlopen", side_effect=error):
+ with pytest.raises(LoginError, match="Verification failed: HTTP 500"):
+ _verify_and_get_user_info("token")
+
+ def test_network_error_raises_login_error(self) -> None:
+ error = URLError(reason="Connection refused")
+ with patch("urllib.request.urlopen", side_effect=error):
+ with pytest.raises(LoginError, match="Could not connect to platform"):
+ _verify_and_get_user_info("token")
+
+ def test_invalid_json_raises_login_error(self) -> None:
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = b"not json"
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ with pytest.raises(LoginError, match="Invalid response from platform"):
+ _verify_and_get_user_info("token")
+
+ def test_naive_expires_at_raises_login_error(self) -> None:
+ """A naive (non-timezone-aware) expires_at should trigger a LoginError."""
+ naive_iso = datetime.now().isoformat() # no tz info
+ body = _make_verify_response(expires_at=naive_iso)
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ with pytest.raises(LoginError, match="expected timezone-aware"):
+ _verify_and_get_user_info("token")
+
+ def test_missing_user_fields_fallback_to_defaults(self) -> None:
+ """Missing user/org fields should fall back to empty strings."""
+ expires_str = (datetime.now(timezone.utc) + timedelta(days=30)).isoformat()
+ body = _make_verify_response(
+ user={},
+ organization={},
+ expires_at=expires_str,
+ )
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ user, org, _, _ = _verify_and_get_user_info("token")
+
+ assert user.id == ""
+ assert user.email == ""
+ assert user.name is None
+ assert org.id == ""
+ assert org.name == ""
+ assert org.role == "member"
+
+ def test_no_token_id_in_response(self) -> None:
+ expires_str = (datetime.now(timezone.utc) + timedelta(days=30)).isoformat()
+ body = _make_verify_response(
+ expires_at=expires_str,
+ token_id=None,
+ )
+ # Remove token_id from payload
+ data = json.loads(body)
+ data.pop("token_id", None)
+ body = json.dumps(data).encode()
+
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+
+ with patch("urllib.request.urlopen", return_value=mock_resp):
+ _, _, _, tid = _verify_and_get_user_info("token")
+
+ assert tid is None
+
+
+# ---------------------------------------------------------------------------
+# login() full flow
+# ---------------------------------------------------------------------------
+
+
+class TestLogin:
+ """Tests for the top-level login() orchestrator."""
+
+ def _stub_verify_response(self) -> MagicMock:
+ """Build a mock urlopen response for successful verification."""
+ expires_str = (datetime.now(timezone.utc) + timedelta(days=90)).isoformat()
+ body = _make_verify_response(expires_at=expires_str)
+ mock_resp = MagicMock()
+ mock_resp.read.return_value = body
+ mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+ mock_resp.__exit__ = MagicMock(return_value=False)
+ return mock_resp
+
+ def test_no_available_port_raises_login_error(self) -> None:
+ with patch("osmosis_ai.auth.flow.find_available_port", return_value=None):
+ with pytest.raises(LoginError, match="No available port"):
+ login()
+
+ def test_successful_login_flow(self) -> None:
+ """Full happy path: port found -> browser opens -> callback -> verify -> save."""
+ mock_verify_resp = self._stub_verify_response()
+
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("test-token-123", None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser") as mock_wb,
+ patch("urllib.request.urlopen", return_value=mock_verify_resp),
+ patch("osmosis_ai.auth.flow.save_credentials") as mock_save,
+ ):
+ mock_wb.open.return_value = True
+ result = login()
+
+ assert isinstance(result, LoginResult)
+ assert result.user.email == "u@test.com"
+ assert result.organization.name == "TestOrg"
+ mock_save.assert_called_once()
+
+ def test_no_browser_mode_does_not_open_browser(self) -> None:
+ mock_verify_resp = self._stub_verify_response()
+
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("test-token", None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser") as mock_wb,
+ patch("urllib.request.urlopen", return_value=mock_verify_resp),
+ patch("osmosis_ai.auth.flow.save_credentials"),
+ ):
+ login(no_browser=True)
+
+ mock_wb.open.assert_not_called()
+
+ def test_callback_returns_error(self) -> None:
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = (None, "User denied")
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = False
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser"),
+ ):
+ with pytest.raises(LoginError, match="Authentication failed: User denied"):
+ login()
+
+ def test_callback_returns_no_token(self) -> None:
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = (None, None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = False
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser"),
+ ):
+ with pytest.raises(LoginError, match="No token received"):
+ login()
+
+ def test_verification_failure_propagates(self) -> None:
+ """If _verify_and_get_user_info fails, login() should propagate LoginError."""
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("bad-token", None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ verify_error = HTTPError(
+ url="http://test",
+ code=401,
+ msg="Unauthorized",
+ hdrs=None,
+ fp=None, # type: ignore[arg-type]
+ )
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser"),
+ patch("urllib.request.urlopen", side_effect=verify_error),
+ ):
+ with pytest.raises(LoginError, match="Invalid or expired token"):
+ login()
+
+ # set_verification_result should have been called with success=False
+ mock_server.set_verification_result.assert_called_once()
+ call_kwargs = mock_server.set_verification_result.call_args
+ assert call_kwargs[1].get("success") is False or call_kwargs[0][0] is False
+
+ def test_server_shutdown_called_in_finally(self) -> None:
+ """The server should always be cleaned up even on failure."""
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = (None, "error")
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = False
+ mock_server._shutdown_event = MagicMock()
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser"),
+ ):
+ with pytest.raises(LoginError):
+ login()
+
+ mock_server.server_close.assert_called_once()
+
+ def test_browser_open_failure_does_not_raise(self) -> None:
+ """If webbrowser.open returns False, login should still proceed."""
+ mock_verify_resp = self._stub_verify_response()
+
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("test-token", None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser") as mock_wb,
+ patch("urllib.request.urlopen", return_value=mock_verify_resp),
+ patch("osmosis_ai.auth.flow.save_credentials"),
+ ):
+ mock_wb.open.return_value = False
+ result = login()
+
+ assert isinstance(result, LoginResult)
+
+ def test_revoked_previous_tokens_count(self) -> None:
+ """revoked_count from server should propagate to LoginResult."""
+ mock_verify_resp = self._stub_verify_response()
+
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("test-token", None)
+ mock_server.revoked_count = 3
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser") as mock_wb,
+ patch("urllib.request.urlopen", return_value=mock_verify_resp),
+ patch("osmosis_ai.auth.flow.save_credentials"),
+ ):
+ mock_wb.open.return_value = True
+ result = login()
+
+ assert result.revoked_previous_tokens == 3
+
+ def test_credentials_saved_with_correct_token(self) -> None:
+ """Credentials passed to save_credentials should contain the actual token."""
+ mock_verify_resp = self._stub_verify_response()
+
+ mock_server = MagicMock()
+ mock_server.wait_for_callback.return_value = ("the-real-token", None)
+ mock_server.revoked_count = 0
+ mock_server._verification_event = MagicMock()
+ mock_server._verification_event.is_set.return_value = True
+
+ with (
+ patch("osmosis_ai.auth.flow.find_available_port", return_value=8976),
+ patch("osmosis_ai.auth.flow.LocalAuthServer", return_value=mock_server),
+ patch("osmosis_ai.auth.flow.webbrowser") as mock_wb,
+ patch("urllib.request.urlopen", return_value=mock_verify_resp),
+ patch("osmosis_ai.auth.flow.save_credentials") as mock_save,
+ ):
+ mock_wb.open.return_value = True
+ login()
+
+ saved_creds: WorkspaceCredentials = mock_save.call_args[0][0]
+ assert saved_creds.access_token == "the-real-token"
+ assert saved_creds.token_type == "Bearer"
diff --git a/tests/test_eval_runner.py b/tests/test_eval_runner.py
index 559b8beb..b4e077e9 100644
--- a/tests/test_eval_runner.py
+++ b/tests/test_eval_runner.py
@@ -879,3 +879,668 @@ def make_baseline() -> MockLLMClient:
)
assert primary_summary.eval_summaries["score_eval"].mean == 1.0
assert baseline_summary.eval_summaries["score_eval"].mean == 0.0
+
+
+# ---------------------------------------------------------------------------
+# Additional coverage: eval function error handling, pass@k edge cases,
+# score aggregation, token statistics, _extract_systemic_error_metrics,
+# _compute_summaries edge cases, _filter_runs_by_tag.
+# ---------------------------------------------------------------------------
+
+from osmosis_ai.rollout.eval.evaluation.runner import (
+ EvalRowResult,
+ EvalRunResult,
+ _extract_systemic_error_metrics,
+)
+
+
+class TestExtractSystemicErrorMetrics:
+ """Tests for the _extract_systemic_error_metrics utility."""
+
+ def test_uses_attributes_when_present(self) -> None:
+ from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
+
+ e = SystemicProviderError("test")
+ e.duration_ms = 1234.5 # type: ignore[attr-defined]
+ e.tokens = 42 # type: ignore[attr-defined]
+ dur, tok = _extract_systemic_error_metrics(e, fallback_started_at=0.0)
+ assert dur == 1234.5
+ assert tok == 42
+
+ def test_fallback_duration_when_missing(self) -> None:
+ import time
+
+ from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
+
+ e = SystemicProviderError("test")
+ # No duration_ms attribute
+ started = time.monotonic() - 0.5 # 500ms ago
+ dur, tok = _extract_systemic_error_metrics(e, fallback_started_at=started)
+ # Duration should be >= 500ms (within reason)
+ assert dur >= 400.0
+ assert tok == 0
+
+ def test_fallback_tokens_when_none(self) -> None:
+ from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
+
+ e = SystemicProviderError("test")
+ e.duration_ms = 100.0 # type: ignore[attr-defined]
+ e.tokens = None # type: ignore[attr-defined]
+ dur, tok = _extract_systemic_error_metrics(e, fallback_started_at=0.0)
+ assert dur == 100.0
+ assert tok == 0
+
+
+class TestEvalFunctionErrorHandling:
+ """Tests for eval function exception and type error handling."""
+
+ @pytest.mark.asyncio
+ async def test_eval_fn_exception_returns_zero_score(self) -> None:
+ """When an eval function raises, its score should be 0.0."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def failing_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ raise ValueError("boom")
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(failing_eval, "failing_eval")],
+ )
+
+ result = await runner.run_single(
+ row=create_sample_row(0),
+ row_index=0,
+ run_index=0,
+ )
+
+ assert result.success is True
+ assert result.scores["failing_eval"] == 0.0
+
+ @pytest.mark.asyncio
+ async def test_eval_fn_returning_non_float_returns_zero(self) -> None:
+ """When an eval function returns a non-numeric type, it should be caught
+ by EvalFnWrapper and the score should default to 0.0 (via the exception
+ handler in run_single)."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def bad_return_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return "not_a_float" # type: ignore[return-value]
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(bad_return_eval, "bad_return_eval")],
+ )
+
+ result = await runner.run_single(
+ row=create_sample_row(0),
+ row_index=0,
+ run_index=0,
+ )
+
+ assert result.success is True
+ # EvalFnWrapper raises EvalFnError which is caught by run_single -> 0.0
+ assert result.scores["bad_return_eval"] == 0.0
+
+ @pytest.mark.asyncio
+ async def test_multiple_eval_fns_one_fails(self) -> None:
+ """If one eval fn fails, others should still be scored correctly."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def good_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 0.75
+
+ def bad_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ raise RuntimeError("eval exploded")
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[
+ EvalFnWrapper(good_eval, "good_eval"),
+ EvalFnWrapper(bad_eval, "bad_eval"),
+ ],
+ )
+
+ result = await runner.run_single(
+ row=create_sample_row(0),
+ row_index=0,
+ run_index=0,
+ )
+
+ assert result.success is True
+ assert result.scores["good_eval"] == 0.75
+ assert result.scores["bad_eval"] == 0.0
+
+
+class TestPassAtKEdgeCases:
+ """Edge cases for pass@k computation in _compute_summaries."""
+
+ @pytest.mark.asyncio
+ async def test_all_runs_pass(self) -> None:
+ """When all runs pass, pass@k should be 1.0 for all k."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def perfect_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(perfect_eval, "perfect_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(i) for i in range(3)],
+ n_runs=5,
+ pass_threshold=1.0,
+ )
+
+ summary = eval_result.eval_summaries["perfect_eval"]
+ assert summary.mean == 1.0
+ assert summary.min == 1.0
+ assert summary.max == 1.0
+ # pass@1 through pass@5 should all be 1.0
+ for k in [1, 3, 5]:
+ assert summary.pass_at_k[k] == 1.0
+
+ @pytest.mark.asyncio
+ async def test_no_runs_pass(self) -> None:
+ """When no runs pass, pass@k should be 0.0 for all k."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def zero_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 0.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(zero_eval, "zero_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0)],
+ n_runs=5,
+ pass_threshold=0.5,
+ )
+
+ summary = eval_result.eval_summaries["zero_eval"]
+ assert summary.mean == 0.0
+ for k in [1, 3, 5]:
+ assert summary.pass_at_k[k] == 0.0
+
+ @pytest.mark.asyncio
+ async def test_pass_at_k_not_computed_for_single_run(self) -> None:
+ """pass@k should not be computed when n_runs == 1."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0)],
+ n_runs=1,
+ )
+
+ summary = eval_result.eval_summaries["simple_eval"]
+ assert summary.pass_at_k == {}
+
+ @pytest.mark.asyncio
+ async def test_pass_at_k_skips_k_greater_than_n_runs(self) -> None:
+ """pass@k should not be computed for k > n_runs."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0)],
+ n_runs=2,
+ pass_threshold=0.5,
+ )
+
+ summary = eval_result.eval_summaries["simple_eval"]
+ # k=1 should be present, k=3, 5, 10 should not
+ assert 1 in summary.pass_at_k
+ assert 3 not in summary.pass_at_k
+ assert 5 not in summary.pass_at_k
+ assert 10 not in summary.pass_at_k
+
+
+class TestScoreAggregation:
+ """Tests for score aggregation across multiple runs and rows."""
+
+ @pytest.mark.asyncio
+ async def test_mean_std_with_varied_scores(self) -> None:
+ """Verify mean and std computation with varying scores."""
+ import math
+
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ scores = iter([0.0, 0.5, 1.0, 0.5])
+
+ def varying_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return next(scores)
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(varying_eval, "varying_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0), create_sample_row(1)],
+ n_runs=2,
+ pass_threshold=0.5,
+ )
+
+ summary = eval_result.eval_summaries["varying_eval"]
+ # Scores: [0.0, 0.5, 1.0, 0.5] -> mean=0.5
+ assert summary.mean == pytest.approx(0.5)
+ assert summary.min == 0.0
+ assert summary.max == 1.0
+ # Population std: sqrt(((0-0.5)^2+(0.5-0.5)^2+(1-0.5)^2+(0.5-0.5)^2)/4)
+ expected_std = math.sqrt((0.25 + 0.0 + 0.25 + 0.0) / 4)
+ assert summary.std == pytest.approx(expected_std)
+
+ @pytest.mark.asyncio
+ async def test_progress_callback_invoked(self) -> None:
+ """on_progress callback should be called for each run."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ progress_calls: list[tuple[int, int]] = []
+
+ def on_progress(current: int, total: int, _result: EvalRunResult) -> None:
+ progress_calls.append((current, total))
+
+ await runner.run_eval(
+ rows=[create_sample_row(0), create_sample_row(1)],
+ n_runs=2,
+ on_progress=on_progress,
+ )
+
+ assert len(progress_calls) == 4
+ assert progress_calls[-1] == (4, 4)
+
+ @pytest.mark.asyncio
+ async def test_start_index_offset(self) -> None:
+ """start_index should offset row_index in results."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0), create_sample_row(1)],
+ n_runs=1,
+ start_index=10,
+ )
+
+ assert eval_result.rows[0].row_index == 10
+ assert eval_result.rows[1].row_index == 11
+
+
+class TestTokenStatistics:
+ """Tests for token usage accumulation."""
+
+ @pytest.mark.asyncio
+ async def test_tokens_accumulated_across_runs(self) -> None:
+ """Total tokens should be sum of all runs."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0), create_sample_row(1)],
+ n_runs=3,
+ )
+
+ # MockLLMClient: each chat call uses 10 prompt + 5 completion = 15 tokens
+ # 2 rows * 3 runs * 1 chat call each = 6 runs * 15 tokens = 90 total
+ assert eval_result.total_tokens == 90
+ assert eval_result.total_runs == 6
+
+ @pytest.mark.asyncio
+ async def test_failed_runs_contribute_zero_tokens(self) -> None:
+ """Failed agent runs that never call LLM should have 0 tokens."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(
+ tools=[create_sample_tool()],
+ run_error=RuntimeError("fail"),
+ )
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0)],
+ n_runs=2,
+ )
+
+ assert eval_result.total_tokens == 0
+ for row in eval_result.rows:
+ for run in row.runs:
+ assert run.tokens == 0
+
+ @pytest.mark.asyncio
+ async def test_duration_is_positive(self) -> None:
+ """All runs should have positive duration."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ eval_result = await runner.run_eval(
+ rows=[create_sample_row(0)],
+ n_runs=1,
+ )
+
+ assert eval_result.total_duration_ms > 0
+ for row in eval_result.rows:
+ for run in row.runs:
+ assert run.duration_ms > 0
+
+
+class TestComputeSummariesEdgeCases:
+ """Tests for _compute_summaries with edge cases."""
+
+ @pytest.mark.asyncio
+ async def test_empty_rows_produces_empty_summaries(self) -> None:
+ """_compute_summaries with empty row results should produce zero summaries."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()])
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ summaries = runner._compute_summaries([], n_runs=1, pass_threshold=1.0)
+ assert "simple_eval" in summaries
+ summary = summaries["simple_eval"]
+ assert summary.mean == 0.0
+ assert summary.std == 0.0
+ assert summary.min == 0.0
+ assert summary.max == 0.0
+
+ @pytest.mark.asyncio
+ async def test_missing_score_defaults_to_zero(self) -> None:
+ """Runs with missing eval score should count as 0.0 in aggregation."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()])
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ )
+
+ # Construct row results where some runs have scores and some don't
+ row_results = [
+ EvalRowResult(
+ row_index=0,
+ runs=[
+ EvalRunResult(
+ run_index=0,
+ success=True,
+ scores={"simple_eval": 1.0},
+ tokens=10,
+ ),
+ EvalRunResult(
+ run_index=1,
+ success=False,
+ scores={}, # missing score - should be 0.0
+ tokens=0,
+ ),
+ ],
+ ),
+ ]
+
+ summaries = runner._compute_summaries(row_results, n_runs=2, pass_threshold=0.5)
+ summary = summaries["simple_eval"]
+ assert summary.mean == pytest.approx(0.5)
+ assert summary.min == 0.0
+ assert summary.max == 1.0
+
+
+class TestFilterRunsByTag:
+ """Tests for _filter_runs_by_tag."""
+
+ def test_filters_correctly(self) -> None:
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()])
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[],
+ )
+
+ row_results = [
+ EvalRowResult(
+ row_index=0,
+ runs=[
+ EvalRunResult(run_index=0, success=True, model_tag="primary"),
+ EvalRunResult(run_index=0, success=True, model_tag="baseline"),
+ ],
+ ),
+ EvalRowResult(
+ row_index=1,
+ runs=[
+ EvalRunResult(run_index=0, success=True, model_tag="primary"),
+ EvalRunResult(run_index=0, success=True, model_tag="baseline"),
+ ],
+ ),
+ ]
+
+ primary_only = runner._filter_runs_by_tag(row_results, "primary")
+ assert len(primary_only) == 2
+ for row in primary_only:
+ assert all(r.model_tag == "primary" for r in row.runs)
+
+ def test_no_matching_tag_returns_empty(self) -> None:
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()])
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[],
+ )
+
+ row_results = [
+ EvalRowResult(
+ row_index=0,
+ runs=[
+ EvalRunResult(run_index=0, success=True, model_tag="primary"),
+ ],
+ ),
+ ]
+
+ result = runner._filter_runs_by_tag(row_results, "baseline")
+ assert result == []
+
+
+class TestRunSingleRouting:
+ """Tests for runner routing based on model_tag."""
+
+ @pytest.mark.asyncio
+ async def test_baseline_tag_uses_baseline_runner(self) -> None:
+ """model_tag='baseline' should route to the baseline rollout runner."""
+ primary_client = MockLLMClient(model="primary-model")
+ baseline_client = MockLLMClient(model="baseline-model")
+ agent = MockAgentLoop(tools=[create_sample_tool()], call_llm=True)
+
+ def simple_eval(
+ solution_str: str,
+ ground_truth: str,
+ extra_info: dict[str, Any],
+ ) -> float:
+ return 1.0
+
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=primary_client, # type: ignore[arg-type]
+ eval_fns=[EvalFnWrapper(simple_eval, "simple_eval")],
+ baseline_llm_client=baseline_client, # type: ignore[arg-type]
+ )
+
+ assert runner.has_baseline is True
+
+ primary_result = await runner.run_single(
+ row=create_sample_row(0),
+ row_index=0,
+ run_index=0,
+ model_tag="primary",
+ )
+ assert primary_result.model_tag == "primary"
+ assert primary_result.success is True
+
+ baseline_result = await runner.run_single(
+ row=create_sample_row(0),
+ row_index=0,
+ run_index=0,
+ model_tag="baseline",
+ )
+ assert baseline_result.model_tag == "baseline"
+ assert baseline_result.success is True
+
+ @pytest.mark.asyncio
+ async def test_no_baseline_has_baseline_false(self) -> None:
+ """Without baseline_llm_client, has_baseline should be False."""
+ client = MockLLMClient()
+ agent = MockAgentLoop(tools=[create_sample_tool()])
+ runner = EvalRunner(
+ agent_loop=agent,
+ llm_client=client, # type: ignore[arg-type]
+ eval_fns=[],
+ )
+ assert runner.has_baseline is False
diff --git a/tests/test_rollout_base.py b/tests/test_rollout_base.py
index 52e24f02..a6f1f6b2 100644
--- a/tests/test_rollout_base.py
+++ b/tests/test_rollout_base.py
@@ -274,8 +274,10 @@ async def test_rollout_context_chat_calls_llm(
def test_agent_loop_subclass_requires_name() -> None:
- """Verify subclass without name raises TypeError."""
- with pytest.raises(TypeError, match="must define a 'name' class attribute"):
+ """Verify subclass without name raises TypeError mentioning the class name."""
+ with pytest.raises(
+ TypeError, match=r"NoNameLoop.*must define a 'name' class attribute"
+ ):
class NoNameLoop(RolloutAgentLoop):
def get_tools(self, request):
@@ -328,8 +330,9 @@ def get_tools(self, request):
# Missing run() method
- with pytest.raises(TypeError, match="abstract"):
+ with pytest.raises(TypeError, match="run") as exc_info:
PartialLoop()
+ assert "abstract" in str(exc_info.value).lower()
@pytest.mark.asyncio
diff --git a/tests/test_rollout_exceptions.py b/tests/test_rollout_exceptions.py
index 19d8e995..1b26993d 100644
--- a/tests/test_rollout_exceptions.py
+++ b/tests/test_rollout_exceptions.py
@@ -16,6 +16,8 @@
from __future__ import annotations
+import pytest
+
from osmosis_ai.rollout import (
AgentLoopNotFoundError,
OsmosisRolloutError,
@@ -26,13 +28,6 @@
)
-def test_osmosis_rollout_error_is_base_exception() -> None:
- """Verify OsmosisRolloutError is the base class."""
- error = OsmosisRolloutError("test error")
- assert isinstance(error, Exception)
- assert str(error) == "test error"
-
-
def test_osmosis_transport_error() -> None:
"""Verify OsmosisTransportError inherits from base."""
error = OsmosisTransportError("connection failed")
@@ -48,11 +43,12 @@ def test_osmosis_server_error_with_status_code() -> None:
assert str(error) == "internal server error"
-def test_osmosis_server_error_different_status_codes() -> None:
- """Verify OsmosisServerError works with various 5xx codes."""
- for status_code in [500, 502, 503, 504]:
- error = OsmosisServerError(f"error {status_code}", status_code)
- assert error.status_code == status_code
+@pytest.mark.parametrize("status_code", [400, 404, 422, 500, 502, 503])
+def test_osmosis_server_error_different_status_codes(status_code: int) -> None:
+ """Verify OsmosisServerError works with various status codes."""
+ error = OsmosisServerError(f"Error {status_code}", status_code)
+ assert error.status_code == status_code
+ assert str(status_code) in str(error)
def test_osmosis_validation_error_with_status_code() -> None:
@@ -63,11 +59,12 @@ def test_osmosis_validation_error_with_status_code() -> None:
assert str(error) == "bad request"
-def test_osmosis_validation_error_different_status_codes() -> None:
- """Verify OsmosisValidationError works with various 4xx codes."""
- for status_code in [400, 401, 403, 404, 422]:
- error = OsmosisValidationError(f"error {status_code}", status_code)
- assert error.status_code == status_code
+@pytest.mark.parametrize("status_code", [400, 422])
+def test_osmosis_validation_error_different_status_codes(status_code: int) -> None:
+ """Verify OsmosisValidationError works with various status codes."""
+ error = OsmosisValidationError(f"Validation {status_code}", status_code)
+ assert error.status_code == status_code
+ assert str(status_code) in str(error)
def test_osmosis_timeout_error() -> None:
@@ -100,15 +97,6 @@ def test_agent_loop_not_found_error_empty_available() -> None:
assert "[]" in str(error)
-def test_exception_inheritance_hierarchy() -> None:
- """Verify all exceptions inherit from OsmosisRolloutError."""
- assert issubclass(OsmosisTransportError, OsmosisRolloutError)
- assert issubclass(OsmosisServerError, OsmosisRolloutError)
- assert issubclass(OsmosisValidationError, OsmosisRolloutError)
- assert issubclass(OsmosisTimeoutError, OsmosisRolloutError)
- assert issubclass(AgentLoopNotFoundError, OsmosisRolloutError)
-
-
def test_exceptions_can_be_caught_by_base() -> None:
"""Verify all exceptions can be caught by base class."""
exceptions = [
diff --git a/tests/test_rollout_schemas.py b/tests/test_rollout_schemas.py
index fff9a68a..b4d7a3a3 100644
--- a/tests/test_rollout_schemas.py
+++ b/tests/test_rollout_schemas.py
@@ -20,7 +20,6 @@
from pydantic import ValidationError
from osmosis_ai.rollout import (
- DEFAULT_MAX_METADATA_SIZE_BYTES,
CompletionsRequest,
CompletionUsage,
InitResponse,
@@ -56,22 +55,12 @@ def test_rollout_request_valid() -> None:
assert request.max_tokens_total == 8192 # default
-def test_rollout_request_empty_rollout_id_rejected() -> None:
- """Verify empty rollout_id is rejected."""
- with pytest.raises(ValidationError, match="rollout_id"):
+@pytest.mark.parametrize("invalid_id", ["", " ", "\t\n"])
+def test_rollout_request_invalid_rollout_id_rejected(invalid_id: str) -> None:
+ """Verify empty or whitespace-only rollout_id is rejected."""
+ with pytest.raises(ValidationError):
RolloutRequest(
- rollout_id="",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
-
-def test_rollout_request_whitespace_rollout_id_rejected() -> None:
- """Verify whitespace-only rollout_id is rejected."""
- with pytest.raises(ValidationError, match="cannot be empty or whitespace"):
- RolloutRequest(
- rollout_id=" ",
+ rollout_id=invalid_id,
server_url="http://localhost:8080",
messages=[],
completion_params={},
@@ -136,18 +125,11 @@ def test_max_metadata_size_configurable() -> None:
set_max_metadata_size_bytes(original_size)
-def test_max_metadata_size_default_value() -> None:
- """Verify default metadata size is 1MB."""
- assert DEFAULT_MAX_METADATA_SIZE_BYTES == 1 * 1024 * 1024
-
-
-def test_set_max_metadata_size_rejects_non_positive() -> None:
+@pytest.mark.parametrize("invalid_size", [0, -1, -100])
+def test_set_max_metadata_size_rejects_non_positive(invalid_size: int) -> None:
"""Verify set_max_metadata_size_bytes rejects non-positive values."""
with pytest.raises(ValueError, match="must be positive"):
- set_max_metadata_size_bytes(0)
-
- with pytest.raises(ValueError, match="must be positive"):
- set_max_metadata_size_bytes(-1)
+ set_max_metadata_size_bytes(invalid_size)
def test_rollout_request_metadata_valid() -> None:
@@ -162,21 +144,6 @@ def test_rollout_request_metadata_valid() -> None:
assert request.metadata["key"] == "value"
-def test_rollout_request_default_values() -> None:
- """Verify default values are set correctly."""
- request = RolloutRequest(
- rollout_id="test-123",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
- assert request.max_turns == 10
- assert request.max_tokens_total == 8192
- assert request.metadata == {}
- assert request.api_key is None
- assert request.tool_server_url is None
-
-
# =============================================================================
# InitResponse Tests
# =============================================================================
@@ -200,12 +167,6 @@ def test_init_response_empty_tools() -> None:
assert response.tools == []
-def test_init_response_default_tools() -> None:
- """Verify InitResponse defaults to empty tools."""
- response = InitResponse(rollout_id="test-123")
- assert response.tools == []
-
-
# =============================================================================
# RolloutResponse Tests
# =============================================================================
@@ -312,19 +273,6 @@ def test_completions_request_valid() -> None:
assert request.logprobs is True
-def test_completions_request_default_values() -> None:
- """Verify CompletionsRequest default values."""
- request = CompletionsRequest(
- rollout_id="test-123",
- messages=[],
- )
- assert request.temperature == 1.0
- assert request.top_p == 1.0
- assert request.max_tokens == 512
- assert request.stop is None
- assert request.logprobs is True
-
-
def test_completions_request_rollout_id_validation() -> None:
"""Verify CompletionsRequest validates rollout_id."""
with pytest.raises(ValidationError, match="cannot be empty or whitespace"):
@@ -357,19 +305,6 @@ def test_completions_request_custom_params() -> None:
# =============================================================================
-def test_rollout_metrics_default_values() -> None:
- """Verify RolloutMetrics default values are zero."""
- metrics = RolloutMetrics()
- assert metrics.total_latency_ms == 0.0
- assert metrics.llm_latency_ms == 0.0
- assert metrics.tool_latency_ms == 0.0
- assert metrics.num_llm_calls == 0
- assert metrics.num_tool_calls == 0
- assert metrics.prompt_tokens == 0
- assert metrics.response_tokens == 0
- assert metrics.max_context_tokens == 0
-
-
def test_rollout_metrics_all_fields() -> None:
"""Verify RolloutMetrics with all fields set."""
metrics = RolloutMetrics(
@@ -397,14 +332,6 @@ def test_rollout_metrics_all_fields() -> None:
# =============================================================================
-def test_completion_usage_defaults() -> None:
- """Verify CompletionUsage default values."""
- usage = CompletionUsage()
- assert usage.prompt_tokens == 0
- assert usage.completion_tokens == 0
- assert usage.total_tokens == 0
-
-
def test_completion_usage_all_fields() -> None:
"""Verify CompletionUsage with all fields set."""
usage = CompletionUsage(
@@ -415,21 +342,3 @@ def test_completion_usage_all_fields() -> None:
assert usage.prompt_tokens == 100
assert usage.completion_tokens == 50
assert usage.total_tokens == 150
-
-
-# =============================================================================
-# RolloutStatus Tests
-# =============================================================================
-
-
-def test_rollout_status_values() -> None:
- """Verify RolloutStatus enum values."""
- assert RolloutStatus.COMPLETED.value == "COMPLETED"
- assert RolloutStatus.ERROR.value == "ERROR"
-
-
-def test_rollout_status_is_string_enum() -> None:
- """Verify RolloutStatus can be compared with strings."""
- # RolloutStatus inherits from str, so can compare directly
- assert RolloutStatus.COMPLETED == "COMPLETED"
- assert RolloutStatus.ERROR == "ERROR"
diff --git a/tests/test_rollout_server.py b/tests/test_rollout_server.py
index 7974cc0a..8030e291 100644
--- a/tests/test_rollout_server.py
+++ b/tests/test_rollout_server.py
@@ -31,6 +31,7 @@
create_app,
)
from osmosis_ai.rollout.server import AppState
+from tests.conftest import FailingAgentLoop, SimpleAgentLoop
# Import FastAPI test client
try:
@@ -42,33 +43,6 @@
FASTAPI_AVAILABLE = False
-class SimpleAgentLoop(RolloutAgentLoop):
- """Simple agent loop for testing."""
-
- name = "simple_agent"
-
- def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
- self._tools = tools or []
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return self._tools
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- return ctx.complete(list(ctx.request.messages))
-
-
-class FailingAgentLoop(RolloutAgentLoop):
- """Agent loop that raises an error."""
-
- name = "failing_agent"
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- raise RuntimeError("Test error")
-
-
# =============================================================================
# AppState Tests
# =============================================================================
@@ -322,6 +296,11 @@ def test_init_endpoint_returns_202() -> None:
)
assert response.status_code == 202
+ data = response.json()
+ assert "tools" in data
+ assert "rollout_id" in data
+ assert data["rollout_id"] == "test-123"
+ assert isinstance(data["tools"], list)
@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
@@ -704,20 +683,6 @@ def test_init_endpoint_without_api_key_configured() -> None:
assert response.status_code == 202
-# =============================================================================
-# Import Error Test
-# =============================================================================
-
-
-def test_create_app_raises_without_fastapi() -> None:
- """Verify create_app raises ImportError when FastAPI not available."""
- # This test only makes sense if we mock the import
- with patch.dict("sys.modules", {"fastapi": None}):
- # We need to reload the module to trigger the import error
- # For this test, we just verify the function exists
- pass # Skip actual test as it requires complex module manipulation
-
-
# =============================================================================
# Debug Logging Tests
# =============================================================================
diff --git a/tests/test_rollout_server_app.py b/tests/test_rollout_server_app.py
index 26180e59..0c5bed87 100644
--- a/tests/test_rollout_server_app.py
+++ b/tests/test_rollout_server_app.py
@@ -22,12 +22,17 @@
OpenAIFunctionToolSchema,
RolloutAgentLoop,
RolloutContext,
- RolloutMetrics,
RolloutRequest,
RolloutResult,
create_app,
)
from osmosis_ai.rollout.server.app import _extract_bearer_token
+from tests.conftest import (
+ SimpleAgentLoop,
+ SlowAgentLoop,
+ make_rollout_payload,
+ mock_llm_client,
+)
# Import FastAPI / httpx test utilities
try:
@@ -42,41 +47,10 @@
# =============================================================================
-# Helpers — reusable agent loop stubs
+# Helpers — module-specific agent loop stubs
# =============================================================================
-class SimpleAgentLoop(RolloutAgentLoop):
- """Agent loop that completes immediately."""
-
- name = "simple_agent"
-
- def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
- self._tools = tools or []
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return self._tools
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- return ctx.complete(list(ctx.request.messages))
-
-
-class SlowAgentLoop(RolloutAgentLoop):
- """Agent loop that sleeps for a configurable duration before completing."""
-
- name = "slow_agent"
-
- def __init__(self, delay: float = 0.5):
- self._delay = delay
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- await asyncio.sleep(self._delay)
- return ctx.complete(list(ctx.request.messages))
-
-
class FailingGetToolsAgentLoop(RolloutAgentLoop):
"""Agent loop whose get_tools() raises an exception."""
@@ -89,34 +63,6 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
return ctx.complete([])
-def _make_rollout_payload(
- rollout_id: str = "test-123",
- idempotency_key: str | None = None,
- **overrides: Any,
-) -> dict[str, Any]:
- """Build a minimal valid RolloutRequest JSON payload."""
- payload: dict[str, Any] = {
- "rollout_id": rollout_id,
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- }
- if idempotency_key is not None:
- payload["idempotency_key"] = idempotency_key
- payload.update(overrides)
- return payload
-
-
-def _mock_llm_client() -> MagicMock:
- """Create a mocked OsmosisLLMClient usable as an async context manager."""
- mock = AsyncMock()
- mock.__aenter__ = AsyncMock(return_value=mock)
- mock.__aexit__ = AsyncMock(return_value=None)
- mock.complete_rollout = AsyncMock()
- mock.get_metrics = MagicMock(return_value=RolloutMetrics())
- return mock
-
-
# =============================================================================
# _extract_bearer_token() tests
# =============================================================================
@@ -354,7 +300,7 @@ async def test_semaphore_limits_concurrent_rollouts(self) -> None:
agent = SlowAgentLoop(delay=0.3)
app = create_app(agent, max_concurrent=max_concurrent)
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -366,7 +312,7 @@ async def test_semaphore_limits_concurrent_rollouts(self) -> None:
for i in range(4):
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id=f"conc-{i}"),
+ json=make_rollout_payload(rollout_id=f"conc-{i}"),
)
assert resp.status_code == 202
@@ -402,7 +348,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
return ctx.complete(list(ctx.request.messages))
app = create_app(OrderTrackingAgent(), max_concurrent=1)
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -412,11 +358,11 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
) as client:
await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="r1"),
+ json=make_rollout_payload(rollout_id="r1"),
)
await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="r2"),
+ json=make_rollout_payload(rollout_id="r2"),
)
# Wait for both to finish
await asyncio.sleep(0.5)
@@ -445,7 +391,7 @@ async def test_idempotency_key_used_when_provided(self) -> None:
"""Two requests with different rollout_ids but same idempotency_key
should be treated as duplicates."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -455,13 +401,13 @@ async def test_idempotency_key_used_when_provided(self) -> None:
) as client:
resp1 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="id-1", idempotency_key="shared-key"
),
)
resp2 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="id-2", idempotency_key="shared-key"
),
)
@@ -474,7 +420,7 @@ async def test_idempotency_key_used_when_provided(self) -> None:
async def test_different_idempotency_keys_are_independent(self) -> None:
"""Requests with different idempotency_keys should be treated separately."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -484,13 +430,13 @@ async def test_different_idempotency_keys_are_independent(self) -> None:
) as client:
resp1 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="id-a", idempotency_key="key-a"
),
)
resp2 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="id-b", idempotency_key="key-b"
),
)
@@ -502,7 +448,7 @@ async def test_different_idempotency_keys_are_independent(self) -> None:
async def test_fallback_to_rollout_id_when_no_idempotency_key(self) -> None:
"""Without idempotency_key, rollout_id is the idempotency key."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -512,11 +458,11 @@ async def test_fallback_to_rollout_id_when_no_idempotency_key(self) -> None:
) as client:
resp1 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="dup-id"),
+ json=make_rollout_payload(rollout_id="dup-id"),
)
resp2 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="dup-id"),
+ json=make_rollout_payload(rollout_id="dup-id"),
)
assert resp1.status_code == 202
@@ -541,7 +487,7 @@ async def test_get_tools_error_returns_500(self) -> None:
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="fail-tools"),
+ json=make_rollout_payload(rollout_id="fail-tools"),
)
assert resp.status_code == 500
@@ -567,7 +513,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
return ctx.complete([])
app = create_app(FailOnceThenSucceedAgent())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -578,14 +524,14 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
# First request fails
resp1 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="retry-me"),
+ json=make_rollout_payload(rollout_id="retry-me"),
)
assert resp1.status_code == 500
# Retry should succeed (init record was cleared)
resp2 = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="retry-me"),
+ json=make_rollout_payload(rollout_id="retry-me"),
)
assert resp2.status_code == 202
@@ -606,7 +552,7 @@ async def test_infrastructure_error_marks_completed(self) -> None:
) as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="infra-err"),
+ json=make_rollout_payload(rollout_id="infra-err"),
)
assert resp.status_code == 202
@@ -630,7 +576,7 @@ class TestBackgroundRolloutTask:
async def test_successful_rollout_calls_complete_rollout(self) -> None:
"""A successful agent run should call llm.complete_rollout with COMPLETED."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -640,7 +586,7 @@ async def test_successful_rollout_calls_complete_rollout(self) -> None:
) as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="bg-ok"),
+ json=make_rollout_payload(rollout_id="bg-ok"),
)
assert resp.status_code == 202
await asyncio.sleep(0.2)
@@ -665,7 +611,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
raise RuntimeError("Agent crashed")
app = create_app(ErrorAgent())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -675,7 +621,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
) as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="bg-err"),
+ json=make_rollout_payload(rollout_id="bg-err"),
)
assert resp.status_code == 202
await asyncio.sleep(0.2)
@@ -688,7 +634,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
async def test_rollout_uses_request_api_key_for_callback(self) -> None:
"""The OsmosisLLMClient should be constructed with the request's api_key."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -698,7 +644,7 @@ async def test_rollout_uses_request_api_key_for_callback(self) -> None:
) as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="api-key-cb", api_key="callback-secret"
),
)
@@ -712,7 +658,7 @@ async def test_rollout_uses_request_api_key_for_callback(self) -> None:
async def test_rollout_passes_server_url_to_client(self) -> None:
"""The OsmosisLLMClient should receive the server_url from the request."""
app = create_app(SimpleAgentLoop())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -722,7 +668,7 @@ async def test_rollout_passes_server_url_to_client(self) -> None:
) as client:
resp = await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(
+ json=make_rollout_payload(
rollout_id="url-test",
server_url="http://traingate:9000",
),
@@ -755,17 +701,6 @@ def test_raw_token_without_bearer_prefix_accepted(self) -> None:
)
assert resp.status_code == 200
- def test_empty_authorization_header_returns_401(self) -> None:
- """An empty Authorization header should be rejected."""
- app = create_app(SimpleAgentLoop(), api_key="some-key")
-
- with TestClient(app) as client:
- resp = client.get(
- "/platform/health",
- headers={"Authorization": ""},
- )
- assert resp.status_code == 401
-
def test_platform_health_returns_agent_name_and_counts(self) -> None:
"""Authenticated /platform/health should return agent info and counts."""
api_key = "check-key"
@@ -801,41 +736,7 @@ def test_raw_token_accepted_on_init(self) -> None:
resp = client.post(
"/v1/rollout/init",
headers={"Authorization": api_key},
- json=_make_rollout_payload(rollout_id="raw-auth"),
- )
- assert resp.status_code == 202
-
- def test_empty_authorization_header_rejected_on_init(self) -> None:
- """An empty Authorization header should be rejected on init."""
- app = create_app(SimpleAgentLoop(), api_key="some-key")
-
- with TestClient(app) as client:
- resp = client.post(
- "/v1/rollout/init",
- headers={"Authorization": ""},
- json=_make_rollout_payload(rollout_id="empty-auth"),
- )
- assert resp.status_code == 401
-
- def test_no_auth_header_at_all_rejected(self) -> None:
- """Missing Authorization header should be rejected when api_key is set."""
- app = create_app(SimpleAgentLoop(), api_key="some-key")
-
- with TestClient(app) as client:
- resp = client.post(
- "/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="no-auth"),
- )
- assert resp.status_code == 401
-
- def test_no_api_key_configured_allows_unauthenticated(self) -> None:
- """When api_key is not set, requests without auth should succeed."""
- app = create_app(SimpleAgentLoop())
-
- with TestClient(app) as client:
- resp = client.post(
- "/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="no-auth-needed"),
+ json=make_rollout_payload(rollout_id="raw-auth"),
)
assert resp.status_code == 202
@@ -852,7 +753,7 @@ async def test_health_returns_correct_counts(self) -> None:
"""Health endpoint should report accurate active and completed counts."""
agent = SlowAgentLoop(delay=0.3)
app = create_app(agent)
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -869,7 +770,7 @@ async def test_health_returns_correct_counts(self) -> None:
# Start a rollout
await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="count-1"),
+ json=make_rollout_payload(rollout_id="count-1"),
)
await asyncio.sleep(0.05)
@@ -884,22 +785,6 @@ async def test_health_returns_correct_counts(self) -> None:
data = health.json()
assert data["completed_rollouts"] >= 1
- def test_health_includes_agent_loop_name(self) -> None:
- """Health response should contain the agent loop name."""
- app = create_app(SimpleAgentLoop())
- with TestClient(app) as client:
- resp = client.get("/health")
- data = resp.json()
- assert data["agent_loop"] == "simple_agent"
-
- def test_health_returns_status_healthy(self) -> None:
- """Health endpoint should return status 'healthy'."""
- app = create_app(SimpleAgentLoop())
- with TestClient(app) as client:
- resp = client.get("/health")
- assert resp.status_code == 200
- assert resp.json()["status"] == "healthy"
-
# =============================================================================
# create_app configuration variations
@@ -971,7 +856,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
return ctx.complete([])
app = create_app(InspectingAgent(), debug_dir=debug_dir)
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -981,7 +866,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
) as client:
await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="debug-prop"),
+ json=make_rollout_payload(rollout_id="debug-prop"),
)
await asyncio.sleep(0.2)
@@ -1005,7 +890,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
return ctx.complete([])
app = create_app(InspectingAgent())
- mock_client = _mock_llm_client()
+ mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
MockCls.return_value = mock_client
@@ -1015,7 +900,7 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
) as client:
await client.post(
"/v1/rollout/init",
- json=_make_rollout_payload(rollout_id="no-debug"),
+ json=make_rollout_payload(rollout_id="no-debug"),
)
await asyncio.sleep(0.2)
diff --git a/tests/test_rollout_state.py b/tests/test_rollout_state.py
new file mode 100644
index 00000000..e4624a1c
--- /dev/null
+++ b/tests/test_rollout_state.py
@@ -0,0 +1,599 @@
+# Copyright 2025 Osmosis AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for osmosis_ai.rollout.server.state (AppState)."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from unittest.mock import MagicMock, patch
+
+from osmosis_ai.rollout.config.settings import RolloutServerSettings
+from osmosis_ai.rollout.server.state import AppState
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def _make_state(
+ max_concurrent: int = 10,
+ record_ttl_seconds: float = 60.0,
+ cleanup_interval_seconds: float = 300.0,
+) -> AppState:
+ """Create an AppState with explicit settings to avoid loading env vars."""
+ settings = RolloutServerSettings(
+ max_concurrent_rollouts=max_concurrent,
+ record_ttl_seconds=record_ttl_seconds,
+ cleanup_interval_seconds=cleanup_interval_seconds,
+ )
+ return AppState(
+ max_concurrent=max_concurrent,
+ record_ttl_seconds=record_ttl_seconds,
+ cleanup_interval_seconds=cleanup_interval_seconds,
+ settings=settings,
+ )
+
+
+# =============================================================================
+# Initialization Tests
+# =============================================================================
+
+
+class TestAppStateInit:
+ """Tests for AppState initialization."""
+
+ def test_initial_state_empty(self) -> None:
+ """AppState starts with empty dicts and no cleanup task."""
+ state = _make_state()
+ assert state.rollout_tasks == {}
+ assert state.completed_rollouts == {}
+ assert state._init_futures == {}
+ assert state._cleanup_task is None
+
+ def test_active_count_starts_at_zero(self) -> None:
+ """active_count property returns 0 when nothing is tracked."""
+ state = _make_state()
+ assert state.active_count == 0
+
+ def test_completed_count_starts_at_zero(self) -> None:
+ """completed_count property returns 0 when nothing is tracked."""
+ state = _make_state()
+ assert state.completed_count == 0
+
+ def test_custom_max_concurrent(self) -> None:
+ """max_concurrent is stored from constructor arg."""
+ state = _make_state(max_concurrent=42)
+ assert state._max_concurrent == 42
+
+ def test_custom_record_ttl(self) -> None:
+ """record_ttl is stored from constructor arg."""
+ state = _make_state(record_ttl_seconds=120.0)
+ assert state.record_ttl == 120.0
+
+ def test_custom_cleanup_interval(self) -> None:
+ """cleanup_interval is stored from constructor arg."""
+ state = _make_state(cleanup_interval_seconds=30.0)
+ assert state._cleanup_interval == 30.0
+
+ def test_semaphore_limits_concurrency(self) -> None:
+ """Semaphore is created with the max_concurrent value."""
+ state = _make_state(max_concurrent=5)
+ # asyncio.Semaphore stores the initial value as _value
+ assert state.semaphore._value == 5
+
+
+# =============================================================================
+# get_or_create_init_future Tests
+# =============================================================================
+
+
+class TestGetOrCreateInitFuture:
+ """Tests for AppState.get_or_create_init_future (deduplication)."""
+
+ async def test_first_call_creates_future(self) -> None:
+ """First call for a key creates a new future and returns created=True."""
+ state = _make_state()
+ future, created = state.get_or_create_init_future("rollout-1")
+ assert created is True
+ assert isinstance(future, asyncio.Future)
+
+ async def test_second_call_returns_same_future(self) -> None:
+ """Second call with same key returns the same future with created=False."""
+ state = _make_state()
+ future1, created1 = state.get_or_create_init_future("rollout-1")
+ future2, created2 = state.get_or_create_init_future("rollout-1")
+
+ assert created1 is True
+ assert created2 is False
+ assert future1 is future2
+
+ async def test_different_keys_get_different_futures(self) -> None:
+ """Different rollout_ids get independent futures."""
+ state = _make_state()
+ future_a, created_a = state.get_or_create_init_future("rollout-a")
+ future_b, created_b = state.get_or_create_init_future("rollout-b")
+
+ assert created_a is True
+ assert created_b is True
+ assert future_a is not future_b
+
+ async def test_future_is_stored_internally(self) -> None:
+ """The created future is stored in _init_futures dict."""
+ state = _make_state()
+ future, _ = state.get_or_create_init_future("rollout-x")
+ assert "rollout-x" in state._init_futures
+ assert state._init_futures["rollout-x"] is future
+
+
+# =============================================================================
+# clear_init_record Tests
+# =============================================================================
+
+
+class TestClearInitRecord:
+ """Tests for AppState.clear_init_record."""
+
+ async def test_clear_existing_record(self) -> None:
+ """Clearing an existing init future removes it."""
+ state = _make_state()
+ state.get_or_create_init_future("rollout-1")
+ assert "rollout-1" in state._init_futures
+
+ state.clear_init_record("rollout-1")
+ assert "rollout-1" not in state._init_futures
+
+ def test_clear_nonexistent_record_no_error(self) -> None:
+ """Clearing a non-existent key does not raise."""
+ state = _make_state()
+ # Should not raise
+ state.clear_init_record("nonexistent")
+
+ async def test_clear_only_affects_target_key(self) -> None:
+ """Clearing one key leaves other keys intact."""
+ state = _make_state()
+ state.get_or_create_init_future("rollout-a")
+ state.get_or_create_init_future("rollout-b")
+
+ state.clear_init_record("rollout-a")
+ assert "rollout-a" not in state._init_futures
+ assert "rollout-b" in state._init_futures
+
+
+# =============================================================================
+# is_duplicate Tests
+# =============================================================================
+
+
+class TestIsDuplicate:
+ """Tests for AppState.is_duplicate (three conditions)."""
+
+ async def test_duplicate_when_active_future_exists(self) -> None:
+ """Returns True when rollout_id has an active init future."""
+ state = _make_state()
+ state.get_or_create_init_future("rollout-1")
+ assert state.is_duplicate("rollout-1") is True
+
+ def test_duplicate_when_in_completed_history(self) -> None:
+ """Returns True when rollout_id is in completed_rollouts."""
+ state = _make_state()
+ state.completed_rollouts["rollout-1"] = time.monotonic()
+ assert state.is_duplicate("rollout-1") is True
+
+ def test_duplicate_when_in_rollout_tasks(self) -> None:
+ """Returns True when rollout_id has an active task."""
+ state = _make_state()
+ mock_task = MagicMock()
+ state.rollout_tasks["rollout-1"] = mock_task # type: ignore[assignment]
+ assert state.is_duplicate("rollout-1") is True
+
+ def test_not_duplicate_for_unknown_key(self) -> None:
+ """Returns False for a completely unknown rollout_id."""
+ state = _make_state()
+ assert state.is_duplicate("unknown") is False
+
+ async def test_not_duplicate_after_clearing_all_records(self) -> None:
+ """After clearing init record and no task/completed, not duplicate."""
+ state = _make_state()
+ state.get_or_create_init_future("rollout-1")
+ state.clear_init_record("rollout-1")
+ assert state.is_duplicate("rollout-1") is False
+
+
+# =============================================================================
+# mark_started / mark_completed Tests
+# =============================================================================
+
+
+class TestMarkStartedAndCompleted:
+ """Tests for AppState.mark_started and mark_completed."""
+
+ def test_mark_started_records_task(self) -> None:
+ """mark_started stores the task in rollout_tasks."""
+ state = _make_state()
+ mock_task = MagicMock()
+ state.mark_started("rollout-1", mock_task) # type: ignore[arg-type]
+ assert "rollout-1" in state.rollout_tasks
+ assert state.rollout_tasks["rollout-1"] is mock_task
+
+ def test_mark_started_increments_active_count(self) -> None:
+ """active_count increases after mark_started."""
+ state = _make_state()
+ assert state.active_count == 0
+ state.mark_started("r1", MagicMock()) # type: ignore[arg-type]
+ assert state.active_count == 1
+ state.mark_started("r2", MagicMock()) # type: ignore[arg-type]
+ assert state.active_count == 2
+
+ def test_mark_completed_removes_from_active(self) -> None:
+ """mark_completed removes the task from rollout_tasks."""
+ state = _make_state()
+ state.mark_started("rollout-1", MagicMock()) # type: ignore[arg-type]
+ state.mark_completed("rollout-1")
+ assert "rollout-1" not in state.rollout_tasks
+
+ def test_mark_completed_adds_to_completed_history(self) -> None:
+ """mark_completed adds entry to completed_rollouts."""
+ state = _make_state()
+ state.mark_started("rollout-1", MagicMock()) # type: ignore[arg-type]
+ state.mark_completed("rollout-1")
+ assert "rollout-1" in state.completed_rollouts
+ assert state.completed_count == 1
+
+ def test_mark_completed_records_monotonic_time(self) -> None:
+ """mark_completed records a monotonic timestamp."""
+ state = _make_state()
+ before = time.monotonic()
+ state.mark_completed("rollout-1")
+ after = time.monotonic()
+ ts = state.completed_rollouts["rollout-1"]
+ assert before <= ts <= after
+
+ def test_mark_completed_nonexistent_key_safe(self) -> None:
+ """mark_completed for a key not in rollout_tasks doesn't raise."""
+ state = _make_state()
+ # Should not raise
+ state.mark_completed("nonexistent")
+ assert "nonexistent" in state.completed_rollouts
+
+ def test_full_lifecycle(self) -> None:
+ """Full lifecycle: start -> complete -> is_duplicate checks."""
+ state = _make_state()
+ mock_task = MagicMock()
+
+ # Before start: not a duplicate
+ assert state.is_duplicate("rollout-1") is False
+
+ # After start: is duplicate (in rollout_tasks)
+ state.mark_started("rollout-1", mock_task) # type: ignore[arg-type]
+ assert state.is_duplicate("rollout-1") is True
+ assert state.active_count == 1
+
+ # After complete: still duplicate (in completed_rollouts)
+ state.mark_completed("rollout-1")
+ assert state.is_duplicate("rollout-1") is True
+ assert state.active_count == 0
+ assert state.completed_count == 1
+
+
+# =============================================================================
+# _prune_completed_records Tests (TTL expiry)
+# =============================================================================
+
+
+class TestPruneCompletedRecords:
+ """Tests for AppState._prune_completed_records."""
+
+ def test_expired_records_are_pruned(self) -> None:
+ """Records older than TTL are removed by pruning."""
+ state = _make_state(record_ttl_seconds=60.0)
+
+ # Insert a record with a timestamp well in the past
+ state.completed_rollouts["old"] = time.monotonic() - 120.0
+ assert state.completed_count == 1
+
+ state._prune_completed_records()
+ assert "old" not in state.completed_rollouts
+ assert state.completed_count == 0
+
+ def test_recent_records_are_kept(self) -> None:
+ """Records within TTL are not pruned."""
+ state = _make_state(record_ttl_seconds=60.0)
+
+ state.completed_rollouts["recent"] = time.monotonic()
+ state._prune_completed_records()
+ assert "recent" in state.completed_rollouts
+
+ def test_mixed_expired_and_recent(self) -> None:
+ """Only expired records are pruned; recent ones remain."""
+ state = _make_state(record_ttl_seconds=60.0)
+
+ now = time.monotonic()
+ state.completed_rollouts["old1"] = now - 120.0
+ state.completed_rollouts["old2"] = now - 90.0
+ state.completed_rollouts["recent1"] = now - 10.0
+ state.completed_rollouts["recent2"] = now
+
+ state._prune_completed_records()
+
+ assert "old1" not in state.completed_rollouts
+ assert "old2" not in state.completed_rollouts
+ assert "recent1" in state.completed_rollouts
+ assert "recent2" in state.completed_rollouts
+
+ async def test_prune_also_removes_init_futures(self) -> None:
+ """Pruning expired completed records also removes their init futures."""
+ state = _make_state(record_ttl_seconds=60.0)
+
+ state.get_or_create_init_future("old-rollout")
+ state.completed_rollouts["old-rollout"] = time.monotonic() - 120.0
+
+ state._prune_completed_records()
+
+ assert "old-rollout" not in state.completed_rollouts
+ assert "old-rollout" not in state._init_futures
+
+ def test_prune_empty_state_no_error(self) -> None:
+ """Pruning with no completed records does not raise."""
+ state = _make_state()
+ state._prune_completed_records()
+ assert state.completed_count == 0
+
+ def test_prune_with_time_mock(self) -> None:
+ """Test pruning with mocked monotonic time for precision."""
+ state = _make_state(record_ttl_seconds=100.0)
+
+ # Insert a record at "time 1000"
+ state.completed_rollouts["r1"] = 1000.0
+ state.completed_rollouts["r2"] = 1050.0
+
+ # Mock time to be 1110 (r1 expired at TTL=100, r2 still valid)
+ with patch(
+ "osmosis_ai.rollout.server.state.time.monotonic", return_value=1110.0
+ ):
+ state._prune_completed_records()
+
+ assert "r1" not in state.completed_rollouts
+ assert "r2" in state.completed_rollouts
+
+
+# =============================================================================
+# Cleanup Task Lifecycle Tests
+# =============================================================================
+
+
+class TestCleanupTaskLifecycle:
+ """Tests for start_cleanup_task and stop_cleanup_task."""
+
+ async def test_start_cleanup_creates_task(self) -> None:
+ """start_cleanup_task creates a background asyncio.Task."""
+ state = _make_state(cleanup_interval_seconds=100.0)
+ try:
+ state.start_cleanup_task()
+ assert state._cleanup_task is not None
+ assert isinstance(state._cleanup_task, asyncio.Task)
+ assert not state._cleanup_task.done()
+ finally:
+ await state.stop_cleanup_task()
+
+ async def test_stop_cleanup_cancels_task(self) -> None:
+ """stop_cleanup_task cancels the background task and sets it to None."""
+ state = _make_state(cleanup_interval_seconds=100.0)
+ state.start_cleanup_task()
+ assert state._cleanup_task is not None
+
+ await state.stop_cleanup_task()
+ assert state._cleanup_task is None
+
+ async def test_stop_cleanup_when_not_started(self) -> None:
+ """stop_cleanup_task is safe to call when no task is running."""
+ state = _make_state()
+ # Should not raise
+ await state.stop_cleanup_task()
+ assert state._cleanup_task is None
+
+ async def test_start_cleanup_idempotent(self) -> None:
+ """Calling start_cleanup_task twice doesn't create a second task."""
+ state = _make_state(cleanup_interval_seconds=100.0)
+ try:
+ state.start_cleanup_task()
+ first_task = state._cleanup_task
+
+ state.start_cleanup_task()
+ assert state._cleanup_task is first_task
+ finally:
+ await state.stop_cleanup_task()
+
+ async def test_cleanup_loop_actually_prunes(self) -> None:
+ """Verify the cleanup loop actually calls _prune_completed_records."""
+ state = _make_state(record_ttl_seconds=60.0)
+
+ # Insert an expired record
+ state.completed_rollouts["expired"] = time.monotonic() - 120.0
+
+ # Patch the cleanup interval to be very short so the loop fires fast,
+ # and patch asyncio.sleep inside the module to return immediately.
+ state._cleanup_interval = 0.01
+
+ sleep_call_count = 0
+ original_sleep = asyncio.sleep
+
+ async def fast_sleep(delay):
+ nonlocal sleep_call_count
+ sleep_call_count += 1
+ # Actually sleep a tiny bit to yield control
+ await original_sleep(0.01)
+
+ try:
+ with patch("asyncio.sleep", side_effect=fast_sleep):
+ state.start_cleanup_task()
+ # Give the loop a chance to run at least once
+ await original_sleep(0.05)
+
+ assert "expired" not in state.completed_rollouts
+ finally:
+ await state.stop_cleanup_task()
+
+
+# =============================================================================
+# cancel_all Tests
+# =============================================================================
+
+
+class TestCancelAll:
+ """Tests for AppState.cancel_all."""
+
+ async def test_cancel_all_cancels_pending_tasks(self) -> None:
+ """cancel_all cancels all pending rollout tasks."""
+ state = _make_state()
+
+ async def long_running():
+ await asyncio.sleep(100)
+
+ task1 = asyncio.create_task(long_running())
+ task2 = asyncio.create_task(long_running())
+ state.mark_started("r1", task1)
+ state.mark_started("r2", task2)
+
+ assert state.active_count == 2
+
+ await state.cancel_all()
+
+ assert state.active_count == 0
+ assert state.rollout_tasks == {}
+ assert task1.cancelled()
+ assert task2.cancelled()
+
+ async def test_cancel_all_empty_state(self) -> None:
+ """cancel_all with no tasks does nothing and doesn't raise."""
+ state = _make_state()
+ await state.cancel_all()
+ assert state.active_count == 0
+
+ async def test_cancel_all_clears_rollout_tasks(self) -> None:
+ """cancel_all clears the rollout_tasks dict."""
+ state = _make_state()
+
+ async def long_running():
+ await asyncio.sleep(100)
+
+ state.mark_started("r1", asyncio.create_task(long_running()))
+ await state.cancel_all()
+ assert state.rollout_tasks == {}
+
+ async def test_cancel_all_handles_already_completed_tasks(self) -> None:
+ """cancel_all handles tasks that have already completed."""
+ state = _make_state()
+
+ async def quick_task():
+ return "done"
+
+ task = asyncio.create_task(quick_task())
+ # Let it complete
+ await asyncio.sleep(0.01)
+ assert task.done()
+
+ state.mark_started("r1", task)
+ # Should not raise even though the task is already done
+ await state.cancel_all()
+ assert state.active_count == 0
+
+
+# =============================================================================
+# active_count / completed_count Property Tests
+# =============================================================================
+
+
+class TestCountProperties:
+ """Tests for active_count and completed_count properties."""
+
+ def test_active_count_reflects_rollout_tasks(self) -> None:
+ """active_count matches the number of entries in rollout_tasks."""
+ state = _make_state()
+ state.rollout_tasks["a"] = MagicMock() # type: ignore[assignment]
+ state.rollout_tasks["b"] = MagicMock() # type: ignore[assignment]
+ assert state.active_count == 2
+
+ def test_completed_count_reflects_completed_rollouts(self) -> None:
+ """completed_count matches the number of entries in completed_rollouts."""
+ state = _make_state()
+ state.completed_rollouts["x"] = time.monotonic()
+ state.completed_rollouts["y"] = time.monotonic()
+ state.completed_rollouts["z"] = time.monotonic()
+ assert state.completed_count == 3
+
+ def test_counts_are_independent(self) -> None:
+ """active_count and completed_count track different things."""
+ state = _make_state()
+ state.rollout_tasks["active"] = MagicMock() # type: ignore[assignment]
+ state.completed_rollouts["done"] = time.monotonic()
+ assert state.active_count == 1
+ assert state.completed_count == 1
+
+
+# =============================================================================
+# Integration / Edge-case Tests
+# =============================================================================
+
+
+class TestEdgeCases:
+ """Edge case and integration tests."""
+
+ async def test_concurrent_get_or_create_same_key(self) -> None:
+ """Multiple concurrent calls to get_or_create_init_future with same key."""
+ state = _make_state()
+
+ results = []
+ for _ in range(10):
+ future, created = state.get_or_create_init_future("same-key")
+ results.append((future, created))
+
+ # First should be created, rest should not
+ assert results[0][1] is True
+ for future, created in results[1:]:
+ assert created is False
+ assert future is results[0][0]
+
+ def test_mark_completed_overwrites_previous_completion(self) -> None:
+ """Completing the same key again updates the timestamp."""
+ state = _make_state()
+ state.mark_completed("rollout-1")
+ first_time = state.completed_rollouts["rollout-1"]
+
+ # Small sleep to ensure time difference
+ import time as _time
+
+ _time.sleep(0.001)
+
+ state.mark_completed("rollout-1")
+ second_time = state.completed_rollouts["rollout-1"]
+
+ assert second_time >= first_time
+
+ async def test_restart_cleanup_after_stop(self) -> None:
+ """Can restart cleanup task after stopping it."""
+ state = _make_state(cleanup_interval_seconds=100.0)
+ try:
+ state.start_cleanup_task()
+ first_task = state._cleanup_task
+
+ await state.stop_cleanup_task()
+ assert state._cleanup_task is None
+
+ state.start_cleanup_task()
+ assert state._cleanup_task is not None
+ assert state._cleanup_task is not first_task
+ finally:
+ await state.stop_cleanup_task()
diff --git a/tests/test_rollout_utils.py b/tests/test_rollout_utils.py
index da4d6f51..496848c1 100644
--- a/tests/test_rollout_utils.py
+++ b/tests/test_rollout_utils.py
@@ -18,6 +18,8 @@
from typing import Any
+import pytest
+
from osmosis_ai.rollout import (
count_messages_by_role,
get_message_content,
@@ -78,35 +80,20 @@ def test_parse_tool_calls_empty_list() -> None:
# =============================================================================
-def test_normalize_stop_none() -> None:
- """Verify normalize_stop returns None for None input."""
- assert normalize_stop(None) is None
-
-
-def test_normalize_stop_string() -> None:
- """Verify normalize_stop converts string to list."""
- assert normalize_stop("stop") == ["stop"]
- assert normalize_stop("STOP") == ["STOP"]
-
-
-def test_normalize_stop_list_of_strings() -> None:
- """Verify normalize_stop passes list of strings through."""
- assert normalize_stop(["a", "b"]) == ["a", "b"]
-
-
-def test_normalize_stop_list_converts_to_strings() -> None:
- """Verify normalize_stop converts list items to strings."""
- assert normalize_stop([1, 2, 3]) == ["1", "2", "3"]
-
-
-def test_normalize_stop_empty_string() -> None:
- """Verify normalize_stop handles empty string."""
- assert normalize_stop("") == [""]
-
-
-def test_normalize_stop_empty_list() -> None:
- """Verify normalize_stop handles empty list."""
- assert normalize_stop([]) == []
+@pytest.mark.parametrize(
+ "input_val,expected",
+ [
+ (None, None),
+ ("stop", ["stop"]),
+ (["a", "b"], ["a", "b"]),
+ ([1, 2, 3], ["1", "2", "3"]),
+ ("", [""]),
+ ([], []),
+ ],
+)
+def test_normalize_stop(input_val: Any, expected: Any) -> None:
+ """Verify normalize_stop normalizes various input formats."""
+ assert normalize_stop(input_val) == expected
def test_normalize_stop_invalid_type() -> None:
@@ -155,34 +142,19 @@ def test_get_message_content_not_string() -> None:
# =============================================================================
-def test_get_message_role_user() -> None:
- """Verify get_message_role extracts user role."""
- message = {"role": "user", "content": "Hi"}
- assert get_message_role(message) == "user"
-
-
-def test_get_message_role_assistant() -> None:
- """Verify get_message_role extracts assistant role."""
- message = {"role": "assistant", "content": "Hello"}
- assert get_message_role(message) == "assistant"
-
-
-def test_get_message_role_tool() -> None:
- """Verify get_message_role extracts tool role."""
- message = {"role": "tool", "content": "42", "tool_call_id": "call_123"}
- assert get_message_role(message) == "tool"
-
-
-def test_get_message_role_missing() -> None:
- """Verify get_message_role returns 'unknown' if missing."""
- message = {"content": "Hi"}
- assert get_message_role(message) == "unknown"
-
-
-def test_get_message_role_not_string() -> None:
- """Verify get_message_role returns 'unknown' for non-string role."""
- message = {"role": 123}
- assert get_message_role(message) == "unknown"
+@pytest.mark.parametrize(
+ "message,expected",
+ [
+ ({"role": "user", "content": "Hi"}, "user"),
+ ({"role": "assistant", "content": "Hello"}, "assistant"),
+ ({"role": "tool", "content": "42", "tool_call_id": "call_123"}, "tool"),
+ ({"content": "Hi"}, "unknown"),
+ ({"role": 123}, "unknown"),
+ ],
+)
+def test_get_message_role(message: dict[str, Any], expected: str) -> None:
+ """Verify get_message_role extracts role or returns 'unknown'."""
+ assert get_message_role(message) == expected
# =============================================================================
@@ -190,37 +162,23 @@ def test_get_message_role_not_string() -> None:
# =============================================================================
-def test_is_assistant_message_true() -> None:
- """Verify is_assistant_message returns True for assistant."""
- assert is_assistant_message({"role": "assistant"}) is True
-
-
-def test_is_assistant_message_false() -> None:
- """Verify is_assistant_message returns False for other roles."""
- assert is_assistant_message({"role": "user"}) is False
- assert is_assistant_message({"role": "tool"}) is False
-
-
-def test_is_tool_message_true() -> None:
- """Verify is_tool_message returns True for tool."""
- assert is_tool_message({"role": "tool"}) is True
-
-
-def test_is_tool_message_false() -> None:
- """Verify is_tool_message returns False for other roles."""
- assert is_tool_message({"role": "user"}) is False
- assert is_tool_message({"role": "assistant"}) is False
-
-
-def test_is_user_message_true() -> None:
- """Verify is_user_message returns True for user."""
- assert is_user_message({"role": "user"}) is True
-
-
-def test_is_user_message_false() -> None:
- """Verify is_user_message returns False for other roles."""
- assert is_user_message({"role": "assistant"}) is False
- assert is_user_message({"role": "tool"}) is False
+@pytest.mark.parametrize(
+ "func,role,expected",
+ [
+ (is_assistant_message, "assistant", True),
+ (is_assistant_message, "user", False),
+ (is_tool_message, "tool", True),
+ (is_tool_message, "user", False),
+ (is_user_message, "user", True),
+ (is_user_message, "assistant", False),
+ ],
+)
+def test_is_role_message(func: Any, role: str, expected: bool) -> None:
+ """Verify is_*_message functions return correct bool for given role."""
+ msg: dict[str, Any] = {"role": role, "content": "test"}
+ if role == "tool":
+ msg["tool_call_id"] = "call_123"
+ assert func(msg) == expected
# =============================================================================
diff --git a/tests/test_settings.py b/tests/test_settings.py
new file mode 100644
index 00000000..48f8c2ef
--- /dev/null
+++ b/tests/test_settings.py
@@ -0,0 +1,576 @@
+# Copyright 2025 Osmosis AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for osmosis_ai.rollout.config.settings."""
+
+from __future__ import annotations
+
+import pytest
+from pydantic import ValidationError
+
+from osmosis_ai.rollout.config.settings import (
+ RolloutClientSettings,
+ RolloutServerSettings,
+ RolloutSettings,
+ configure,
+ get_settings,
+ reset_settings,
+)
+
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture(autouse=True)
+def _clean_settings():
+ """Reset the global settings singleton before and after each test."""
+ reset_settings()
+ yield
+ reset_settings()
+
+
+# =============================================================================
+# Default Values Tests
+# =============================================================================
+
+
+class TestDefaultValues:
+ """Verify sensible default values when no env vars are set."""
+
+ def test_client_timeout_seconds_default(self) -> None:
+ """Default client timeout is 300 seconds."""
+ settings = RolloutClientSettings()
+ assert settings.timeout_seconds == 300.0
+
+ def test_client_max_retries_default(self) -> None:
+ """Default max retries is 3."""
+ settings = RolloutClientSettings()
+ assert settings.max_retries == 3
+
+ def test_client_complete_rollout_retries_default(self) -> None:
+ """Default complete_rollout_retries is 2."""
+ settings = RolloutClientSettings()
+ assert settings.complete_rollout_retries == 2
+
+ def test_client_retry_base_delay_default(self) -> None:
+ """Default retry base delay is 1.0."""
+ settings = RolloutClientSettings()
+ assert settings.retry_base_delay == 1.0
+
+ def test_client_retry_max_delay_default(self) -> None:
+ """Default retry max delay is 30.0."""
+ settings = RolloutClientSettings()
+ assert settings.retry_max_delay == 30.0
+
+ def test_client_max_connections_default(self) -> None:
+ """Default max connections is 100."""
+ settings = RolloutClientSettings()
+ assert settings.max_connections == 100
+
+ def test_client_max_keepalive_connections_default(self) -> None:
+ """Default max keepalive connections is 20."""
+ settings = RolloutClientSettings()
+ assert settings.max_keepalive_connections == 20
+
+ def test_server_max_concurrent_rollouts_default(self) -> None:
+ """Default max concurrent rollouts is 100."""
+ settings = RolloutServerSettings()
+ assert settings.max_concurrent_rollouts == 100
+
+ def test_server_record_ttl_seconds_default(self) -> None:
+ """Default record TTL is 3600 seconds (1 hour)."""
+ settings = RolloutServerSettings()
+ assert settings.record_ttl_seconds == 3600.0
+
+ def test_server_cleanup_interval_seconds_default(self) -> None:
+ """Default cleanup interval is 60 seconds."""
+ settings = RolloutServerSettings()
+ assert settings.cleanup_interval_seconds == 60.0
+
+ def test_server_request_timeout_seconds_default(self) -> None:
+ """Default request timeout is 600 seconds."""
+ settings = RolloutServerSettings()
+ assert settings.request_timeout_seconds == 600.0
+
+ def test_server_registration_readiness_timeout_default(self) -> None:
+ """Default registration readiness timeout is 10 seconds."""
+ settings = RolloutServerSettings()
+ assert settings.registration_readiness_timeout_seconds == 10.0
+
+ def test_server_registration_readiness_poll_interval_default(self) -> None:
+ """Default registration readiness poll interval is 0.2 seconds."""
+ settings = RolloutServerSettings()
+ assert settings.registration_readiness_poll_interval_seconds == 0.2
+
+ def test_server_registration_shutdown_timeout_default(self) -> None:
+ """Default registration shutdown timeout is 30 seconds."""
+ settings = RolloutServerSettings()
+ assert settings.registration_shutdown_timeout_seconds == 30.0
+
+ def test_rollout_max_metadata_size_bytes_default(self) -> None:
+ """Default max metadata size is 1MB."""
+ settings = RolloutSettings()
+ assert settings.max_metadata_size_bytes == 1024 * 1024
+
+ def test_rollout_settings_has_client_and_server(self) -> None:
+ """RolloutSettings contains client and server sub-settings."""
+ settings = RolloutSettings()
+ assert isinstance(settings.client, RolloutClientSettings)
+ assert isinstance(settings.server, RolloutServerSettings)
+
+
+# =============================================================================
+# Singleton Pattern Tests (get_settings / configure / reset)
+# =============================================================================
+
+
+class TestSingleton:
+ """Tests for the global settings singleton pattern."""
+
+ def test_get_settings_returns_instance(self) -> None:
+ """get_settings returns a RolloutSettings instance."""
+ settings = get_settings()
+ assert isinstance(settings, RolloutSettings)
+
+ def test_get_settings_returns_same_instance(self) -> None:
+ """Multiple calls to get_settings return the same instance."""
+ s1 = get_settings()
+ s2 = get_settings()
+ assert s1 is s2
+
+ def test_reset_settings_clears_singleton(self) -> None:
+ """reset_settings causes get_settings to create a new instance."""
+ s1 = get_settings()
+ reset_settings()
+ s2 = get_settings()
+ assert s1 is not s2
+
+ def test_configure_sets_custom_settings(self) -> None:
+ """configure() overrides the global singleton."""
+ custom = RolloutSettings(
+ max_metadata_size_bytes=2048,
+ )
+ configure(custom)
+ s = get_settings()
+ assert s is custom
+ assert s.max_metadata_size_bytes == 2048
+
+ def test_configure_overrides_previous_singleton(self) -> None:
+ """configure() replaces any previously set singleton."""
+ s1 = get_settings()
+ custom = RolloutSettings()
+ configure(custom)
+ s2 = get_settings()
+ assert s2 is custom
+ assert s2 is not s1
+
+ def test_reset_after_configure(self) -> None:
+ """reset_settings works after configure."""
+ custom = RolloutSettings(max_metadata_size_bytes=5000)
+ configure(custom)
+ reset_settings()
+ s = get_settings()
+ # Should be a fresh default instance, not the custom one
+ assert s is not custom
+ assert s.max_metadata_size_bytes == 1024 * 1024
+
+
+# =============================================================================
+# configure() Override Tests
+# =============================================================================
+
+
+class TestConfigureOverride:
+ """Tests for programmatic configuration via configure()."""
+
+ def test_configure_with_custom_client(self) -> None:
+ """configure() with custom client settings."""
+ custom = RolloutSettings(
+ client=RolloutClientSettings(timeout_seconds=120.0, max_retries=5),
+ )
+ configure(custom)
+ s = get_settings()
+ assert s.client.timeout_seconds == 120.0
+ assert s.client.max_retries == 5
+
+ def test_configure_with_custom_server(self) -> None:
+ """configure() with custom server settings."""
+ custom = RolloutSettings(
+ server=RolloutServerSettings(max_concurrent_rollouts=500),
+ )
+ configure(custom)
+ s = get_settings()
+ assert s.server.max_concurrent_rollouts == 500
+
+ def test_configure_with_all_custom_values(self) -> None:
+ """configure() with all sub-settings overridden."""
+ custom = RolloutSettings(
+ client=RolloutClientSettings(
+ timeout_seconds=60.0,
+ max_retries=1,
+ retry_base_delay=0.5,
+ ),
+ server=RolloutServerSettings(
+ max_concurrent_rollouts=200,
+ record_ttl_seconds=7200.0,
+ ),
+ max_metadata_size_bytes=2048,
+ )
+ configure(custom)
+ s = get_settings()
+ assert s.client.timeout_seconds == 60.0
+ assert s.client.max_retries == 1
+ assert s.client.retry_base_delay == 0.5
+ assert s.server.max_concurrent_rollouts == 200
+ assert s.server.record_ttl_seconds == 7200.0
+ assert s.max_metadata_size_bytes == 2048
+
+
+# =============================================================================
+# Field Validation Tests
+# =============================================================================
+
+
+class TestFieldValidation:
+ """Tests that invalid values are rejected by Pydantic validation."""
+
+ # --- Client settings validation ---
+
+ def test_client_timeout_too_low(self) -> None:
+ """Client timeout below minimum (1.0) is rejected."""
+ with pytest.raises(ValidationError, match="timeout_seconds"):
+ RolloutClientSettings(timeout_seconds=0.5)
+
+ def test_client_timeout_too_high(self) -> None:
+ """Client timeout above maximum (3600.0) is rejected."""
+ with pytest.raises(ValidationError, match="timeout_seconds"):
+ RolloutClientSettings(timeout_seconds=5000.0)
+
+ def test_client_max_retries_negative(self) -> None:
+ """Negative max_retries is rejected."""
+ with pytest.raises(ValidationError, match="max_retries"):
+ RolloutClientSettings(max_retries=-1)
+
+ def test_client_max_retries_too_high(self) -> None:
+ """max_retries above maximum (10) is rejected."""
+ with pytest.raises(ValidationError, match="max_retries"):
+ RolloutClientSettings(max_retries=20)
+
+ def test_client_retry_base_delay_too_low(self) -> None:
+ """retry_base_delay below minimum (0.1) is rejected."""
+ with pytest.raises(ValidationError, match="retry_base_delay"):
+ RolloutClientSettings(retry_base_delay=0.01)
+
+ def test_client_max_connections_too_low(self) -> None:
+ """max_connections below minimum (1) is rejected."""
+ with pytest.raises(ValidationError, match="max_connections"):
+ RolloutClientSettings(max_connections=0)
+
+ def test_client_max_connections_too_high(self) -> None:
+ """max_connections above maximum (1000) is rejected."""
+ with pytest.raises(ValidationError, match="max_connections"):
+ RolloutClientSettings(max_connections=2000)
+
+ # --- Server settings validation ---
+
+ def test_server_max_concurrent_too_low(self) -> None:
+ """max_concurrent_rollouts below minimum (1) is rejected."""
+ with pytest.raises(ValidationError, match="max_concurrent_rollouts"):
+ RolloutServerSettings(max_concurrent_rollouts=0)
+
+ def test_server_max_concurrent_too_high(self) -> None:
+ """max_concurrent_rollouts above maximum (10000) is rejected."""
+ with pytest.raises(ValidationError, match="max_concurrent_rollouts"):
+ RolloutServerSettings(max_concurrent_rollouts=20000)
+
+ def test_server_record_ttl_too_low(self) -> None:
+ """record_ttl_seconds below minimum (60.0) is rejected."""
+ with pytest.raises(ValidationError, match="record_ttl_seconds"):
+ RolloutServerSettings(record_ttl_seconds=10.0)
+
+ def test_server_record_ttl_too_high(self) -> None:
+ """record_ttl_seconds above maximum (86400.0) is rejected."""
+ with pytest.raises(ValidationError, match="record_ttl_seconds"):
+ RolloutServerSettings(record_ttl_seconds=100000.0)
+
+ def test_server_cleanup_interval_too_low(self) -> None:
+ """cleanup_interval_seconds below minimum (10.0) is rejected."""
+ with pytest.raises(ValidationError, match="cleanup_interval_seconds"):
+ RolloutServerSettings(cleanup_interval_seconds=5.0)
+
+ def test_server_request_timeout_too_low(self) -> None:
+ """request_timeout_seconds below minimum (10.0) is rejected."""
+ with pytest.raises(ValidationError, match="request_timeout_seconds"):
+ RolloutServerSettings(request_timeout_seconds=5.0)
+
+ def test_server_registration_readiness_timeout_too_low(self) -> None:
+ """registration_readiness_timeout_seconds below minimum (1.0) is rejected."""
+ with pytest.raises(
+ ValidationError, match="registration_readiness_timeout_seconds"
+ ):
+ RolloutServerSettings(registration_readiness_timeout_seconds=0.5)
+
+ # --- Global settings validation ---
+
+ def test_max_metadata_size_bytes_too_low(self) -> None:
+ """max_metadata_size_bytes below minimum (1024) is rejected."""
+ with pytest.raises(ValidationError, match="max_metadata_size_bytes"):
+ RolloutSettings(max_metadata_size_bytes=100)
+
+ def test_max_metadata_size_bytes_too_high(self) -> None:
+ """max_metadata_size_bytes above maximum (100MB) is rejected."""
+ with pytest.raises(ValidationError, match="max_metadata_size_bytes"):
+ RolloutSettings(max_metadata_size_bytes=200 * 1024 * 1024)
+
+ # --- Valid boundary values ---
+
+ def test_client_timeout_at_minimum(self) -> None:
+ """Client timeout at exact minimum (1.0) is accepted."""
+ s = RolloutClientSettings(timeout_seconds=1.0)
+ assert s.timeout_seconds == 1.0
+
+ def test_client_timeout_at_maximum(self) -> None:
+ """Client timeout at exact maximum (3600.0) is accepted."""
+ s = RolloutClientSettings(timeout_seconds=3600.0)
+ assert s.timeout_seconds == 3600.0
+
+ def test_server_max_concurrent_at_minimum(self) -> None:
+ """max_concurrent_rollouts at exact minimum (1) is accepted."""
+ s = RolloutServerSettings(max_concurrent_rollouts=1)
+ assert s.max_concurrent_rollouts == 1
+
+ def test_server_max_concurrent_at_maximum(self) -> None:
+ """max_concurrent_rollouts at exact maximum (10000) is accepted."""
+ s = RolloutServerSettings(max_concurrent_rollouts=10000)
+ assert s.max_concurrent_rollouts == 10000
+
+
+# =============================================================================
+# Environment Variable Loading Tests
+# =============================================================================
+
+
+class TestEnvironmentVariables:
+ """Tests for loading settings from environment variables."""
+
+ def test_client_timeout_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Client timeout_seconds loads from OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS", "120.0")
+ settings = RolloutClientSettings()
+ assert settings.timeout_seconds == 120.0
+
+ def test_client_max_retries_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Client max_retries loads from OSMOSIS_ROLLOUT_CLIENT_MAX_RETRIES."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_MAX_RETRIES", "5")
+ settings = RolloutClientSettings()
+ assert settings.max_retries == 5
+
+ def test_client_retry_base_delay_from_env(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Client retry_base_delay loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_RETRY_BASE_DELAY", "2.5")
+ settings = RolloutClientSettings()
+ assert settings.retry_base_delay == 2.5
+
+ def test_client_max_connections_from_env(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Client max_connections loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_MAX_CONNECTIONS", "50")
+ settings = RolloutClientSettings()
+ assert settings.max_connections == 50
+
+ def test_server_max_concurrent_from_env(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Server max_concurrent_rollouts loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS", "200")
+ settings = RolloutServerSettings()
+ assert settings.max_concurrent_rollouts == 200
+
+ def test_server_record_ttl_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Server record_ttl_seconds loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_SERVER_RECORD_TTL_SECONDS", "7200")
+ settings = RolloutServerSettings()
+ assert settings.record_ttl_seconds == 7200.0
+
+ def test_server_cleanup_interval_from_env(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Server cleanup_interval_seconds loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_SERVER_CLEANUP_INTERVAL_SECONDS", "120")
+ settings = RolloutServerSettings()
+ assert settings.cleanup_interval_seconds == 120.0
+
+ def test_server_request_timeout_from_env(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Server request_timeout_seconds loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_SERVER_REQUEST_TIMEOUT_SECONDS", "300")
+ settings = RolloutServerSettings()
+ assert settings.request_timeout_seconds == 300.0
+
+ def test_max_metadata_size_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Global max_metadata_size_bytes loads from env."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES", "2048")
+ settings = RolloutSettings()
+ assert settings.max_metadata_size_bytes == 2048
+
+ def test_invalid_env_var_value_rejected(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Invalid env var value (not a number) is rejected."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS", "not_a_number")
+ with pytest.raises(ValidationError):
+ RolloutClientSettings()
+
+ def test_env_var_violating_constraint_rejected(
+ self, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ """Env var value that violates constraint (too low) is rejected."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS", "0.1")
+ with pytest.raises(ValidationError):
+ RolloutClientSettings()
+
+ def test_get_settings_reads_env_vars(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """get_settings() singleton also picks up env vars on first call."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES", "4096")
+ settings = get_settings()
+ assert settings.max_metadata_size_bytes == 4096
+
+ def test_extra_env_vars_are_ignored(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Unknown env vars with the right prefix are ignored (extra='ignore')."""
+ monkeypatch.setenv("OSMOSIS_ROLLOUT_CLIENT_NONEXISTENT_FIELD", "42")
+ # Should not raise
+ settings = RolloutClientSettings()
+ assert not hasattr(settings, "nonexistent_field")
+
+
+# =============================================================================
+# Programmatic Construction Tests
+# =============================================================================
+
+
+class TestProgrammaticConstruction:
+ """Tests for constructing settings programmatically."""
+
+ def test_construct_client_with_all_fields(self) -> None:
+ """Construct RolloutClientSettings with all fields specified."""
+ s = RolloutClientSettings(
+ timeout_seconds=60.0,
+ max_retries=5,
+ complete_rollout_retries=3,
+ retry_base_delay=2.0,
+ retry_max_delay=60.0,
+ max_connections=200,
+ max_keepalive_connections=50,
+ )
+ assert s.timeout_seconds == 60.0
+ assert s.max_retries == 5
+ assert s.complete_rollout_retries == 3
+ assert s.retry_base_delay == 2.0
+ assert s.retry_max_delay == 60.0
+ assert s.max_connections == 200
+ assert s.max_keepalive_connections == 50
+
+ def test_construct_server_with_all_fields(self) -> None:
+ """Construct RolloutServerSettings with all fields specified."""
+ s = RolloutServerSettings(
+ max_concurrent_rollouts=500,
+ record_ttl_seconds=1800.0,
+ cleanup_interval_seconds=30.0,
+ request_timeout_seconds=120.0,
+ registration_readiness_timeout_seconds=5.0,
+ registration_readiness_poll_interval_seconds=0.5,
+ registration_shutdown_timeout_seconds=60.0,
+ )
+ assert s.max_concurrent_rollouts == 500
+ assert s.record_ttl_seconds == 1800.0
+ assert s.cleanup_interval_seconds == 30.0
+ assert s.request_timeout_seconds == 120.0
+ assert s.registration_readiness_timeout_seconds == 5.0
+ assert s.registration_readiness_poll_interval_seconds == 0.5
+ assert s.registration_shutdown_timeout_seconds == 60.0
+
+ def test_construct_rollout_settings_nested(self) -> None:
+ """Construct RolloutSettings with nested sub-settings."""
+ s = RolloutSettings(
+ client=RolloutClientSettings(timeout_seconds=99.0),
+ server=RolloutServerSettings(max_concurrent_rollouts=42),
+ max_metadata_size_bytes=8192,
+ )
+ assert s.client.timeout_seconds == 99.0
+ assert s.server.max_concurrent_rollouts == 42
+ assert s.max_metadata_size_bytes == 8192
+
+
+# =============================================================================
+# Parametrized Boundary Tests
+# =============================================================================
+
+
+class TestParametrizedBoundaries:
+ """Parametrized tests for boundary validation."""
+
+ @pytest.mark.parametrize(
+ "value",
+ [1.0, 100.0, 3600.0],
+ ids=["min", "mid", "max"],
+ )
+ def test_client_timeout_valid_range(self, value: float) -> None:
+ """Client timeout accepts values in valid range [1.0, 3600.0]."""
+ s = RolloutClientSettings(timeout_seconds=value)
+ assert s.timeout_seconds == value
+
+ @pytest.mark.parametrize(
+ "value",
+ [0.5, 0.0, -1.0, 3601.0],
+ ids=["below_min", "zero", "negative", "above_max"],
+ )
+ def test_client_timeout_invalid_range(self, value: float) -> None:
+ """Client timeout rejects values outside valid range."""
+ with pytest.raises(ValidationError):
+ RolloutClientSettings(timeout_seconds=value)
+
+ @pytest.mark.parametrize(
+ "value",
+ [0, 1, 5, 10],
+ ids=["zero", "one", "five", "ten"],
+ )
+ def test_client_max_retries_valid_range(self, value: int) -> None:
+ """Client max_retries accepts values in valid range [0, 10]."""
+ s = RolloutClientSettings(max_retries=value)
+ assert s.max_retries == value
+
+ @pytest.mark.parametrize(
+ "value",
+ [60.0, 3600.0, 86400.0],
+ ids=["min", "mid", "max"],
+ )
+ def test_server_record_ttl_valid_range(self, value: float) -> None:
+ """Server record_ttl_seconds accepts values in valid range."""
+ s = RolloutServerSettings(record_ttl_seconds=value)
+ assert s.record_ttl_seconds == value
+
+ @pytest.mark.parametrize(
+ "value",
+ [1, 100, 5000, 10000],
+ ids=["min", "default", "mid", "max"],
+ )
+ def test_server_max_concurrent_valid_range(self, value: int) -> None:
+ """Server max_concurrent_rollouts accepts values in valid range."""
+ s = RolloutServerSettings(max_concurrent_rollouts=value)
+ assert s.max_concurrent_rollouts == value
From b71bd50e08f139122f7bff79b9e233b1a98073dc Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 14:30:03 -0800
Subject: [PATCH 15/60] Reorganize tests into `tests/unit/` mirroring source
layout
---
pytest.ini | 3 +++
tests/unit/__init__.py | 0
tests/unit/auth/__init__.py | 0
.../auth/test_credentials.py} | 0
tests/{test_auth_flow.py => unit/auth/test_flow.py} | 0
tests/{ => unit/auth}/test_platform_client.py | 0
tests/unit/cli_services/__init__.py | 0
tests/{ => unit/cli_services}/test_cli_services.py | 0
tests/unit/rollout/__init__.py | 0
tests/unit/rollout/config/__init__.py | 0
tests/{ => unit/rollout/config}/test_settings.py | 0
tests/unit/rollout/core/__init__.py | 0
tests/{test_rollout_base.py => unit/rollout/core/test_base.py} | 0
.../rollout/core/test_exceptions.py} | 0
.../rollout/core/test_schemas.py} | 0
tests/unit/rollout/eval/__init__.py | 0
tests/unit/rollout/eval/common/__init__.py | 0
.../rollout/eval/common/test_dataset.py} | 0
tests/unit/rollout/eval/evaluation/__init__.py | 0
tests/{ => unit/rollout/eval/evaluation}/test_eval_fn.py | 0
.../rollout/eval/evaluation/test_report.py} | 0
.../rollout/eval/evaluation/test_runner.py} | 0
tests/unit/rollout/eval/test_mode/__init__.py | 0
.../rollout/eval/test_mode/test_providers.py} | 0
.../rollout/eval/test_mode/test_runner.py} | 0
tests/unit/rollout/mcp/__init__.py | 0
.../rollout/mcp/test_agent_loop.py} | 0
tests/unit/rollout/server/__init__.py | 0
.../rollout/server/test_api_key.py} | 0
.../rollout/server/test_app.py} | 0
tests/{ => unit/rollout/server}/test_registration.py | 0
.../rollout/server/test_server.py} | 0
.../rollout/server/test_state.py} | 0
tests/{ => unit/rollout}/test_cli_utils.py | 0
tests/{test_rollout_client.py => unit/rollout/test_client.py} | 0
tests/{ => unit/rollout}/test_console.py | 0
tests/{ => unit/rollout}/test_network.py | 0
.../rollout/test_registry.py} | 0
.../{test_rollout_testing.py => unit/rollout/test_testing.py} | 0
tests/{test_rollout_tools.py => unit/rollout/test_tools.py} | 0
tests/{test_rollout_utils.py => unit/rollout/test_utils.py} | 0
.../rollout/test_validator.py} | 0
tests/{ => unit}/test_cli.py | 0
tests/{ => unit}/test_litellm_provider.py | 0
tests/{ => unit}/test_rubric_eval.py | 0
tests/{ => unit}/test_utils_decorators.py | 0
46 files changed, 3 insertions(+)
create mode 100644 tests/unit/__init__.py
create mode 100644 tests/unit/auth/__init__.py
rename tests/{test_auth_credentials.py => unit/auth/test_credentials.py} (100%)
rename tests/{test_auth_flow.py => unit/auth/test_flow.py} (100%)
rename tests/{ => unit/auth}/test_platform_client.py (100%)
create mode 100644 tests/unit/cli_services/__init__.py
rename tests/{ => unit/cli_services}/test_cli_services.py (100%)
create mode 100644 tests/unit/rollout/__init__.py
create mode 100644 tests/unit/rollout/config/__init__.py
rename tests/{ => unit/rollout/config}/test_settings.py (100%)
create mode 100644 tests/unit/rollout/core/__init__.py
rename tests/{test_rollout_base.py => unit/rollout/core/test_base.py} (100%)
rename tests/{test_rollout_exceptions.py => unit/rollout/core/test_exceptions.py} (100%)
rename tests/{test_rollout_schemas.py => unit/rollout/core/test_schemas.py} (100%)
create mode 100644 tests/unit/rollout/eval/__init__.py
create mode 100644 tests/unit/rollout/eval/common/__init__.py
rename tests/{test_test_mode_dataset.py => unit/rollout/eval/common/test_dataset.py} (100%)
create mode 100644 tests/unit/rollout/eval/evaluation/__init__.py
rename tests/{ => unit/rollout/eval/evaluation}/test_eval_fn.py (100%)
rename tests/{test_eval_report.py => unit/rollout/eval/evaluation/test_report.py} (100%)
rename tests/{test_eval_runner.py => unit/rollout/eval/evaluation/test_runner.py} (100%)
create mode 100644 tests/unit/rollout/eval/test_mode/__init__.py
rename tests/{test_test_mode_providers.py => unit/rollout/eval/test_mode/test_providers.py} (100%)
rename tests/{test_test_mode_runner.py => unit/rollout/eval/test_mode/test_runner.py} (100%)
create mode 100644 tests/unit/rollout/mcp/__init__.py
rename tests/{test_mcp_agent_loop.py => unit/rollout/mcp/test_agent_loop.py} (100%)
create mode 100644 tests/unit/rollout/server/__init__.py
rename tests/{test_rollout_api_key.py => unit/rollout/server/test_api_key.py} (100%)
rename tests/{test_rollout_server_app.py => unit/rollout/server/test_app.py} (100%)
rename tests/{ => unit/rollout/server}/test_registration.py (100%)
rename tests/{test_rollout_server.py => unit/rollout/server/test_server.py} (100%)
rename tests/{test_rollout_state.py => unit/rollout/server/test_state.py} (100%)
rename tests/{ => unit/rollout}/test_cli_utils.py (100%)
rename tests/{test_rollout_client.py => unit/rollout/test_client.py} (100%)
rename tests/{ => unit/rollout}/test_console.py (100%)
rename tests/{ => unit/rollout}/test_network.py (100%)
rename tests/{test_rollout_registry.py => unit/rollout/test_registry.py} (100%)
rename tests/{test_rollout_testing.py => unit/rollout/test_testing.py} (100%)
rename tests/{test_rollout_tools.py => unit/rollout/test_tools.py} (100%)
rename tests/{test_rollout_utils.py => unit/rollout/test_utils.py} (100%)
rename tests/{test_rollout_validator.py => unit/rollout/test_validator.py} (100%)
rename tests/{ => unit}/test_cli.py (100%)
rename tests/{ => unit}/test_litellm_provider.py (100%)
rename tests/{ => unit}/test_rubric_eval.py (100%)
rename tests/{ => unit}/test_utils_decorators.py (100%)
diff --git a/pytest.ini b/pytest.ini
index b8b07995..04d3b842 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -5,3 +5,6 @@ filterwarnings =
# Configure asyncio mode for pytest-asyncio
asyncio_mode = auto
+
+# Explicitly define test discovery path
+testpaths = tests
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit/auth/__init__.py b/tests/unit/auth/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_auth_credentials.py b/tests/unit/auth/test_credentials.py
similarity index 100%
rename from tests/test_auth_credentials.py
rename to tests/unit/auth/test_credentials.py
diff --git a/tests/test_auth_flow.py b/tests/unit/auth/test_flow.py
similarity index 100%
rename from tests/test_auth_flow.py
rename to tests/unit/auth/test_flow.py
diff --git a/tests/test_platform_client.py b/tests/unit/auth/test_platform_client.py
similarity index 100%
rename from tests/test_platform_client.py
rename to tests/unit/auth/test_platform_client.py
diff --git a/tests/unit/cli_services/__init__.py b/tests/unit/cli_services/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_cli_services.py b/tests/unit/cli_services/test_cli_services.py
similarity index 100%
rename from tests/test_cli_services.py
rename to tests/unit/cli_services/test_cli_services.py
diff --git a/tests/unit/rollout/__init__.py b/tests/unit/rollout/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit/rollout/config/__init__.py b/tests/unit/rollout/config/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_settings.py b/tests/unit/rollout/config/test_settings.py
similarity index 100%
rename from tests/test_settings.py
rename to tests/unit/rollout/config/test_settings.py
diff --git a/tests/unit/rollout/core/__init__.py b/tests/unit/rollout/core/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_rollout_base.py b/tests/unit/rollout/core/test_base.py
similarity index 100%
rename from tests/test_rollout_base.py
rename to tests/unit/rollout/core/test_base.py
diff --git a/tests/test_rollout_exceptions.py b/tests/unit/rollout/core/test_exceptions.py
similarity index 100%
rename from tests/test_rollout_exceptions.py
rename to tests/unit/rollout/core/test_exceptions.py
diff --git a/tests/test_rollout_schemas.py b/tests/unit/rollout/core/test_schemas.py
similarity index 100%
rename from tests/test_rollout_schemas.py
rename to tests/unit/rollout/core/test_schemas.py
diff --git a/tests/unit/rollout/eval/__init__.py b/tests/unit/rollout/eval/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit/rollout/eval/common/__init__.py b/tests/unit/rollout/eval/common/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_test_mode_dataset.py b/tests/unit/rollout/eval/common/test_dataset.py
similarity index 100%
rename from tests/test_test_mode_dataset.py
rename to tests/unit/rollout/eval/common/test_dataset.py
diff --git a/tests/unit/rollout/eval/evaluation/__init__.py b/tests/unit/rollout/eval/evaluation/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_eval_fn.py b/tests/unit/rollout/eval/evaluation/test_eval_fn.py
similarity index 100%
rename from tests/test_eval_fn.py
rename to tests/unit/rollout/eval/evaluation/test_eval_fn.py
diff --git a/tests/test_eval_report.py b/tests/unit/rollout/eval/evaluation/test_report.py
similarity index 100%
rename from tests/test_eval_report.py
rename to tests/unit/rollout/eval/evaluation/test_report.py
diff --git a/tests/test_eval_runner.py b/tests/unit/rollout/eval/evaluation/test_runner.py
similarity index 100%
rename from tests/test_eval_runner.py
rename to tests/unit/rollout/eval/evaluation/test_runner.py
diff --git a/tests/unit/rollout/eval/test_mode/__init__.py b/tests/unit/rollout/eval/test_mode/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_test_mode_providers.py b/tests/unit/rollout/eval/test_mode/test_providers.py
similarity index 100%
rename from tests/test_test_mode_providers.py
rename to tests/unit/rollout/eval/test_mode/test_providers.py
diff --git a/tests/test_test_mode_runner.py b/tests/unit/rollout/eval/test_mode/test_runner.py
similarity index 100%
rename from tests/test_test_mode_runner.py
rename to tests/unit/rollout/eval/test_mode/test_runner.py
diff --git a/tests/unit/rollout/mcp/__init__.py b/tests/unit/rollout/mcp/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_mcp_agent_loop.py b/tests/unit/rollout/mcp/test_agent_loop.py
similarity index 100%
rename from tests/test_mcp_agent_loop.py
rename to tests/unit/rollout/mcp/test_agent_loop.py
diff --git a/tests/unit/rollout/server/__init__.py b/tests/unit/rollout/server/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_rollout_api_key.py b/tests/unit/rollout/server/test_api_key.py
similarity index 100%
rename from tests/test_rollout_api_key.py
rename to tests/unit/rollout/server/test_api_key.py
diff --git a/tests/test_rollout_server_app.py b/tests/unit/rollout/server/test_app.py
similarity index 100%
rename from tests/test_rollout_server_app.py
rename to tests/unit/rollout/server/test_app.py
diff --git a/tests/test_registration.py b/tests/unit/rollout/server/test_registration.py
similarity index 100%
rename from tests/test_registration.py
rename to tests/unit/rollout/server/test_registration.py
diff --git a/tests/test_rollout_server.py b/tests/unit/rollout/server/test_server.py
similarity index 100%
rename from tests/test_rollout_server.py
rename to tests/unit/rollout/server/test_server.py
diff --git a/tests/test_rollout_state.py b/tests/unit/rollout/server/test_state.py
similarity index 100%
rename from tests/test_rollout_state.py
rename to tests/unit/rollout/server/test_state.py
diff --git a/tests/test_cli_utils.py b/tests/unit/rollout/test_cli_utils.py
similarity index 100%
rename from tests/test_cli_utils.py
rename to tests/unit/rollout/test_cli_utils.py
diff --git a/tests/test_rollout_client.py b/tests/unit/rollout/test_client.py
similarity index 100%
rename from tests/test_rollout_client.py
rename to tests/unit/rollout/test_client.py
diff --git a/tests/test_console.py b/tests/unit/rollout/test_console.py
similarity index 100%
rename from tests/test_console.py
rename to tests/unit/rollout/test_console.py
diff --git a/tests/test_network.py b/tests/unit/rollout/test_network.py
similarity index 100%
rename from tests/test_network.py
rename to tests/unit/rollout/test_network.py
diff --git a/tests/test_rollout_registry.py b/tests/unit/rollout/test_registry.py
similarity index 100%
rename from tests/test_rollout_registry.py
rename to tests/unit/rollout/test_registry.py
diff --git a/tests/test_rollout_testing.py b/tests/unit/rollout/test_testing.py
similarity index 100%
rename from tests/test_rollout_testing.py
rename to tests/unit/rollout/test_testing.py
diff --git a/tests/test_rollout_tools.py b/tests/unit/rollout/test_tools.py
similarity index 100%
rename from tests/test_rollout_tools.py
rename to tests/unit/rollout/test_tools.py
diff --git a/tests/test_rollout_utils.py b/tests/unit/rollout/test_utils.py
similarity index 100%
rename from tests/test_rollout_utils.py
rename to tests/unit/rollout/test_utils.py
diff --git a/tests/test_rollout_validator.py b/tests/unit/rollout/test_validator.py
similarity index 100%
rename from tests/test_rollout_validator.py
rename to tests/unit/rollout/test_validator.py
diff --git a/tests/test_cli.py b/tests/unit/test_cli.py
similarity index 100%
rename from tests/test_cli.py
rename to tests/unit/test_cli.py
diff --git a/tests/test_litellm_provider.py b/tests/unit/test_litellm_provider.py
similarity index 100%
rename from tests/test_litellm_provider.py
rename to tests/unit/test_litellm_provider.py
diff --git a/tests/test_rubric_eval.py b/tests/unit/test_rubric_eval.py
similarity index 100%
rename from tests/test_rubric_eval.py
rename to tests/unit/test_rubric_eval.py
diff --git a/tests/test_utils_decorators.py b/tests/unit/test_utils_decorators.py
similarity index 100%
rename from tests/test_utils_decorators.py
rename to tests/unit/test_utils_decorators.py
From 0903e2c7348662f00db8d2e5162344a2db2b4e12 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 14:41:03 -0800
Subject: [PATCH 16/60] Refactor and expand test suite for osmosis_ai
- Deleted `test_litellm_provider.py` as its contents have been integrated into `test_rubric_eval.py` and `test_utils.py`.
- Added comprehensive tests for prompt building, JSON parsing, and LiteLLM evaluation in `test_rubric_eval.py`.
- Introduced new tests for utility functions in `test_utils.py`, covering decorator validation and type checking.
- Enhanced existing tests for rollout context and external LLM client functionality, ensuring robust coverage and maintainability.
---
tests/unit/rollout/core/test_base.py | 135 +++
.../test_llm_client.py} | 0
tests/unit/rollout/server/test_server.py | 929 ------------------
tests/unit/test_litellm_provider.py | 490 ---------
tests/unit/test_rubric_eval.py | 517 +++++++++-
...test_utils_decorators.py => test_utils.py} | 0
6 files changed, 651 insertions(+), 1420 deletions(-)
rename tests/unit/rollout/eval/{test_mode/test_providers.py => common/test_llm_client.py} (100%)
delete mode 100644 tests/unit/rollout/server/test_server.py
delete mode 100644 tests/unit/test_litellm_provider.py
rename tests/unit/{test_utils_decorators.py => test_utils.py} (100%)
diff --git a/tests/unit/rollout/core/test_base.py b/tests/unit/rollout/core/test_base.py
index a6f1f6b2..6e75ae2c 100644
--- a/tests/unit/rollout/core/test_base.py
+++ b/tests/unit/rollout/core/test_base.py
@@ -394,3 +394,138 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
assert result.status == "ERROR"
assert result.error_message == "Something failed"
+
+
+# =============================================================================
+# RolloutContext Debug Logging Tests
+# =============================================================================
+
+
+def test_rollout_context_debug_disabled_by_default(
+ sample_rollout_request: RolloutRequest,
+) -> None:
+ """Verify debug logging is disabled by default."""
+ mock_llm = MagicMock()
+ ctx = RolloutContext(
+ request=sample_rollout_request,
+ tools=[],
+ llm=mock_llm,
+ )
+
+ assert ctx.debug_enabled is False
+ assert ctx._debug_dir is None
+
+
+def test_rollout_context_debug_enabled_when_dir_set(
+ sample_rollout_request: RolloutRequest,
+) -> None:
+ """Verify debug logging is enabled when _debug_dir is set."""
+ mock_llm = MagicMock()
+ ctx = RolloutContext(
+ request=sample_rollout_request,
+ tools=[],
+ llm=mock_llm,
+ _debug_dir="/tmp/debug",
+ )
+
+ assert ctx.debug_enabled is True
+ assert ctx._debug_dir == "/tmp/debug"
+
+
+def test_rollout_context_log_event_noop_when_disabled(
+ sample_rollout_request: RolloutRequest,
+) -> None:
+ """Verify log_event is a no-op when debug logging is disabled."""
+ mock_llm = MagicMock()
+ ctx = RolloutContext(
+ request=sample_rollout_request,
+ tools=[],
+ llm=mock_llm,
+ )
+
+ # Should not raise or have any effect
+ ctx.log_event("test_event", key="value")
+
+
+def test_rollout_context_log_event_writes_to_file(tmp_path) -> None:
+ """Verify log_event writes events to JSONL file."""
+ import json
+ import os
+
+ debug_dir = str(tmp_path / "debug_logs")
+
+ request = RolloutRequest(
+ rollout_id="test-rollout-456",
+ server_url="http://localhost:8080",
+ messages=[],
+ completion_params={},
+ )
+
+ mock_llm = MagicMock()
+ ctx = RolloutContext(
+ request=request,
+ tools=[],
+ llm=mock_llm,
+ _debug_dir=debug_dir,
+ )
+
+ # Log some events
+ ctx.log_event("pre_llm", turn=0, num_messages=3)
+ ctx.log_event("llm_response", turn=0, has_tool_calls=True)
+ ctx.log_event("rollout_complete", finish_reason="stop", reward=1.0)
+
+ # Verify file was created
+ debug_file = os.path.join(debug_dir, "test-rollout-456.jsonl")
+ assert os.path.exists(debug_file)
+
+ # Verify contents
+ with open(debug_file, encoding="utf-8") as f:
+ lines = f.readlines()
+
+ assert len(lines) == 3
+
+ event1 = json.loads(lines[0])
+ assert event1["event"] == "pre_llm"
+ assert event1["rollout_id"] == "test-rollout-456"
+ assert event1["turn"] == 0
+ assert event1["num_messages"] == 3
+
+ event2 = json.loads(lines[1])
+ assert event2["event"] == "llm_response"
+ assert event2["has_tool_calls"] is True
+
+ event3 = json.loads(lines[2])
+ assert event3["event"] == "rollout_complete"
+ assert event3["finish_reason"] == "stop"
+ assert event3["reward"] == 1.0
+
+
+def test_rollout_context_get_debug_file_path() -> None:
+ """Verify _get_debug_file_path returns correct path."""
+ request = RolloutRequest(
+ rollout_id="my-rollout-id",
+ server_url="http://localhost:8080",
+ messages=[],
+ completion_params={},
+ )
+
+ mock_llm = MagicMock()
+
+ # Without debug_dir
+ ctx_no_debug = RolloutContext(
+ request=request,
+ tools=[],
+ llm=mock_llm,
+ )
+ assert ctx_no_debug._get_debug_file_path() is None
+
+ # With debug_dir
+ ctx_with_debug = RolloutContext(
+ request=request,
+ tools=[],
+ llm=mock_llm,
+ _debug_dir="/var/log/rollouts",
+ )
+ assert (
+ ctx_with_debug._get_debug_file_path() == "/var/log/rollouts/my-rollout-id.jsonl"
+ )
diff --git a/tests/unit/rollout/eval/test_mode/test_providers.py b/tests/unit/rollout/eval/common/test_llm_client.py
similarity index 100%
rename from tests/unit/rollout/eval/test_mode/test_providers.py
rename to tests/unit/rollout/eval/common/test_llm_client.py
diff --git a/tests/unit/rollout/server/test_server.py b/tests/unit/rollout/server/test_server.py
deleted file mode 100644
index 8030e291..00000000
--- a/tests/unit/rollout/server/test_server.py
+++ /dev/null
@@ -1,929 +0,0 @@
-# Copyright 2025 Osmosis AI
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for osmosis_ai.rollout.server."""
-
-from __future__ import annotations
-
-import asyncio
-from typing import Any
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-
-from osmosis_ai.rollout import (
- OpenAIFunctionToolSchema,
- RolloutAgentLoop,
- RolloutContext,
- RolloutRequest,
- RolloutResult,
- create_app,
-)
-from osmosis_ai.rollout.server import AppState
-from tests.conftest import FailingAgentLoop, SimpleAgentLoop
-
-# Import FastAPI test client
-try:
- from fastapi import FastAPI
- from fastapi.testclient import TestClient
-
- FASTAPI_AVAILABLE = True
-except ImportError:
- FASTAPI_AVAILABLE = False
-
-
-# =============================================================================
-# AppState Tests
-# =============================================================================
-
-
-def test_app_state_is_duplicate_not_found() -> None:
- """Verify is_duplicate returns False for unknown rollout."""
- state = AppState()
- assert state.is_duplicate("unknown") is False
-
-
-def test_app_state_is_duplicate_running() -> None:
- """Verify is_duplicate returns True for running rollout."""
- state = AppState()
- mock_task = MagicMock()
- state.mark_started("rollout-1", mock_task)
-
- assert state.is_duplicate("rollout-1") is True
-
-
-def test_app_state_is_duplicate_completed() -> None:
- """Verify is_duplicate returns True for completed rollout."""
- state = AppState()
- state.completed_rollouts["rollout-1"] = 12345.0
-
- assert state.is_duplicate("rollout-1") is True
-
-
-def test_app_state_mark_started_and_completed() -> None:
- """Verify state transitions from started to completed."""
- state = AppState()
- mock_task = MagicMock()
-
- state.mark_started("rollout-1", mock_task)
- assert "rollout-1" in state.rollout_tasks
- assert "rollout-1" not in state.completed_rollouts
-
- state.mark_completed("rollout-1")
- assert "rollout-1" not in state.rollout_tasks
- assert "rollout-1" in state.completed_rollouts
-
-
-def test_app_state_mark_completed_moves_to_completed() -> None:
- """Verify mark_completed removes from tasks and adds to completed."""
- state = AppState()
- mock_task = MagicMock()
-
- state.mark_started("rollout-1", mock_task)
- state.mark_completed("rollout-1")
-
- assert state.rollout_tasks.get("rollout-1") is None
- assert "rollout-1" in state.completed_rollouts
-
-
-def test_app_state_prune_completed_records() -> None:
- """Verify old completed records are pruned."""
- state = AppState(record_ttl_seconds=100)
-
- # Add old record (simulating time in the past)
- state.completed_rollouts["old-rollout"] = 0.0 # Very old timestamp
-
- # Manually trigger prune
- with patch("time.monotonic", return_value=200.0):
- state._prune_completed_records()
-
- assert "old-rollout" not in state.completed_rollouts
-
-
-def test_app_state_prune_keeps_recent_records() -> None:
- """Verify recent completed records are kept."""
- state = AppState(record_ttl_seconds=100)
-
- with patch("time.monotonic", return_value=150.0):
- state.completed_rollouts["recent-rollout"] = 100.0 # 50 seconds old
-
- state._prune_completed_records()
-
- assert "recent-rollout" in state.completed_rollouts
-
-
-@pytest.mark.asyncio
-async def test_app_state_cancel_all() -> None:
- """Verify cancel_all cancels running tasks."""
- state = AppState()
-
- # Create mock tasks
- task1 = AsyncMock()
- task1.cancel = MagicMock()
- task2 = AsyncMock()
- task2.cancel = MagicMock()
-
- state.rollout_tasks["task1"] = task1
- state.rollout_tasks["task2"] = task2
-
- with patch("asyncio.gather", new_callable=AsyncMock):
- await state.cancel_all()
-
- task1.cancel.assert_called_once()
- task2.cancel.assert_called_once()
-
-
-# =============================================================================
-# create_app Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_create_app_returns_fastapi_instance() -> None:
- """Verify create_app returns a FastAPI application."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- assert isinstance(app, FastAPI)
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_create_app_sets_title() -> None:
- """Verify app title includes agent loop name."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- assert "simple_agent" in app.title
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_create_app_with_custom_max_concurrent() -> None:
- """Verify create_app accepts custom max_concurrent."""
- agent_loop = SimpleAgentLoop()
- # Should not raise
- app = create_app(agent_loop, max_concurrent=50)
- assert app is not None
-
-
-# =============================================================================
-# /health Endpoint Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_health_endpoint() -> None:
- """Verify /health returns status information."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.get("/health")
-
- assert response.status_code == 200
- data = response.json()
- assert data["status"] == "healthy"
- assert "active_rollouts" in data
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_health_endpoint_includes_agent_name() -> None:
- """Verify /health includes agent_loop name."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.get("/health")
-
- data = response.json()
- assert data["agent_loop"] == "simple_agent"
-
-
-# =============================================================================
-# /platform/health Endpoint Tests (Authenticated Platform Health Check)
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_platform_health_endpoint_not_exposed_without_api_key() -> None:
- """Verify /platform/health is not exposed when API key auth is disabled."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.get("/platform/health")
-
- assert response.status_code == 404
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_platform_health_endpoint_missing_auth_returns_401() -> None:
- """Verify /platform/health requires Authorization when API key is configured."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop, api_key="test-api-key-12345")
-
- with TestClient(app) as client:
- response = client.get("/platform/health")
-
- assert response.status_code == 401
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_platform_health_endpoint_invalid_auth_returns_401() -> None:
- """Verify /platform/health rejects incorrect API keys."""
- agent_loop = SimpleAgentLoop()
- test_api_key = "test-api-key-12345"
- app = create_app(agent_loop, api_key=test_api_key)
-
- with TestClient(app) as client:
- response = client.get(
- "/platform/health",
- headers={"Authorization": "Bearer wrong-api-key"},
- )
-
- assert response.status_code == 401
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_platform_health_endpoint_valid_auth_returns_200() -> None:
- """Verify /platform/health returns status when given correct API key."""
- agent_loop = SimpleAgentLoop()
- test_api_key = "test-api-key-12345"
- app = create_app(agent_loop, api_key=test_api_key)
-
- with TestClient(app) as client:
- response = client.get(
- "/platform/health",
- headers={"Authorization": f"Bearer {test_api_key}"},
- )
-
- assert response.status_code == 200
- data = response.json()
- assert data["status"] == "healthy"
- assert data["agent_loop"] == "simple_agent"
-
-
-# =============================================================================
-# /v1/rollout/init Endpoint Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_returns_202() -> None:
- """Verify /v1/rollout/init returns 202 Accepted."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "test-123",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
-
- assert response.status_code == 202
- data = response.json()
- assert "tools" in data
- assert "rollout_id" in data
- assert data["rollout_id"] == "test-123"
- assert isinstance(data["tools"], list)
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_returns_tools(
- sample_tool_schema: OpenAIFunctionToolSchema,
-) -> None:
- """Verify /v1/rollout/init returns tools from agent_loop.get_tools()."""
- agent_loop = SimpleAgentLoop(tools=[sample_tool_schema])
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "test-123",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response.status_code == 202
- data = response.json()
- assert data["rollout_id"] == "test-123"
- assert len(data["tools"]) == 1
- assert data["tools"][0]["function"]["name"] == "add"
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_invalid_request() -> None:
- """Verify /v1/rollout/init returns 422 for invalid request."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- # Missing required fields
- "rollout_id": "test-123",
- },
- )
-
- assert response.status_code == 422
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_empty_rollout_id_rejected() -> None:
- """Verify /v1/rollout/init rejects empty rollout_id."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response.status_code == 422
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_whitespace_rollout_id_rejected() -> None:
- """Verify /v1/rollout/init rejects whitespace rollout_id."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": " ",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response.status_code == 422
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_echoes_rollout_id() -> None:
- """Verify /v1/rollout/init echoes back the rollout_id."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "my-unique-id-456",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- data = response.json()
- assert data["rollout_id"] == "my-unique-id-456"
-
-
-# =============================================================================
-# Idempotency Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_idempotent() -> None:
- """Verify duplicate requests return same response without new task."""
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- # First request
- response1 = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "idempotent-test",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- # Second request with same rollout_id
- response2 = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "idempotent-test",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response1.status_code == 202
- assert response2.status_code == 202
- assert response1.json()["rollout_id"] == response2.json()["rollout_id"]
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_idempotent_caches_tools() -> None:
- """Verify get_tools is called once and duplicate requests return the same tools."""
-
- class CountingToolsAgentLoop(RolloutAgentLoop):
- name = "counting_tools_agent"
-
- def __init__(self) -> None:
- self.calls = 0
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- self.calls += 1
- return [
- OpenAIFunctionToolSchema(
- type="function",
- function={
- "name": f"tool_{self.calls}",
- "description": "test tool",
- },
- )
- ]
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- return ctx.complete(list(ctx.request.messages))
-
- agent_loop = CountingToolsAgentLoop()
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response1 = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "idempotent-tools-test",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
- response2 = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "idempotent-tools-test",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response1.status_code == 202
- assert response2.status_code == 202
- assert agent_loop.calls == 1
- assert response1.json()["tools"] == response2.json()["tools"]
-
-
-# =============================================================================
-# Background Task Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-@pytest.mark.asyncio
-async def test_background_rollout_completes() -> None:
- """Verify background task completes and calls complete_rollout."""
- from osmosis_ai.rollout import RolloutMetrics
-
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop)
-
- # Use mock to intercept OsmosisLLMClient
- with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockClient:
- mock_client_instance = AsyncMock()
- mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
- mock_client_instance.__aexit__ = AsyncMock(return_value=None)
- mock_client_instance.complete_rollout = AsyncMock()
- # get_metrics is a sync method, so use MagicMock
- mock_client_instance.get_metrics = MagicMock(return_value=RolloutMetrics())
- MockClient.return_value = mock_client_instance
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "bg-test",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "test"}],
- "completion_params": {},
- # Callback auth: RolloutServer -> TrainGate
- "api_key": "callback-token",
- },
- )
-
- assert response.status_code == 202
-
- # Give background task time to complete
- await asyncio.sleep(0.1)
-
- # Verify complete_rollout was called
- MockClient.assert_called_once()
- assert MockClient.call_args.kwargs["server_url"] == "http://localhost:8080"
- assert MockClient.call_args.kwargs["rollout_id"] == "bg-test"
- assert MockClient.call_args.kwargs["api_key"] == "callback-token"
-
- mock_client_instance.complete_rollout.assert_called_once()
- call_kwargs = mock_client_instance.complete_rollout.call_args[1]
- assert call_kwargs["status"] == "COMPLETED"
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-@pytest.mark.asyncio
-async def test_background_rollout_handles_agent_error() -> None:
- """Verify background task handles agent errors gracefully."""
- agent_loop = FailingAgentLoop()
- app = create_app(agent_loop)
-
- with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockClient:
- mock_client_instance = AsyncMock()
- mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
- mock_client_instance.__aexit__ = AsyncMock(return_value=None)
- mock_client_instance.complete_rollout = AsyncMock()
- MockClient.return_value = mock_client_instance
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "error-test",
- "server_url": "http://localhost:8080",
- "messages": [],
- "completion_params": {},
- },
- )
-
- assert response.status_code == 202
-
- # Give background task time to complete
- await asyncio.sleep(0.1)
-
- # Verify complete_rollout was called with ERROR status
- mock_client_instance.complete_rollout.assert_called()
- call_kwargs = mock_client_instance.complete_rollout.call_args[1]
- assert call_kwargs["status"] == "ERROR"
- assert "Test error" in call_kwargs["error_message"]
-
-
-# =============================================================================
-# API Key Validation Tests
-# =============================================================================
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_with_valid_api_key() -> None:
- """Verify /v1/rollout/init accepts requests with valid API key."""
- agent_loop = SimpleAgentLoop()
- test_api_key = "test-api-key-12345"
- app = create_app(agent_loop, api_key=test_api_key)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- headers={"Authorization": f"Bearer {test_api_key}"},
- json={
- "rollout_id": "api-key-test",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
-
- assert response.status_code == 202
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_with_invalid_api_key() -> None:
- """Verify /v1/rollout/init rejects requests with invalid API key."""
- agent_loop = SimpleAgentLoop()
- test_api_key = "test-api-key-12345"
- app = create_app(agent_loop, api_key=test_api_key)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- headers={"Authorization": "Bearer wrong-api-key"},
- json={
- "rollout_id": "api-key-test",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
-
- assert response.status_code == 401
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_with_missing_api_key() -> None:
- """Verify /v1/rollout/init rejects requests without API key when required."""
- agent_loop = SimpleAgentLoop()
- test_api_key = "test-api-key-12345"
- app = create_app(agent_loop, api_key=test_api_key)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "api-key-test",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- # No api_key provided
- },
- )
-
- assert response.status_code == 401
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_init_endpoint_without_api_key_configured() -> None:
- """Verify /v1/rollout/init works without API key when not configured."""
- agent_loop = SimpleAgentLoop()
- # No api_key parameter - should not require authentication
- app = create_app(agent_loop)
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "no-auth-test",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
-
- assert response.status_code == 202
-
-
-# =============================================================================
-# Debug Logging Tests
-# =============================================================================
-
-
-def test_rollout_context_debug_disabled_by_default() -> None:
- """Verify debug logging is disabled by default."""
- from osmosis_ai.rollout import RolloutRequest
- from osmosis_ai.rollout.core.base import RolloutContext
-
- request = RolloutRequest(
- rollout_id="test-123",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
- mock_llm = MagicMock()
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- )
-
- assert ctx.debug_enabled is False
- assert ctx._debug_dir is None
-
-
-def test_rollout_context_debug_enabled_when_dir_set() -> None:
- """Verify debug logging is enabled when _debug_dir is set."""
- from osmosis_ai.rollout import RolloutRequest
- from osmosis_ai.rollout.core.base import RolloutContext
-
- request = RolloutRequest(
- rollout_id="test-123",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
- mock_llm = MagicMock()
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- _debug_dir="/tmp/debug",
- )
-
- assert ctx.debug_enabled is True
- assert ctx._debug_dir == "/tmp/debug"
-
-
-def test_rollout_context_log_event_noop_when_disabled() -> None:
- """Verify log_event is a no-op when debug logging is disabled."""
- from osmosis_ai.rollout import RolloutRequest
- from osmosis_ai.rollout.core.base import RolloutContext
-
- request = RolloutRequest(
- rollout_id="test-123",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
- mock_llm = MagicMock()
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- )
-
- # Should not raise or have any effect
- ctx.log_event("test_event", key="value")
-
-
-def test_rollout_context_log_event_writes_to_file(tmp_path: Any) -> None:
- """Verify log_event writes events to JSONL file."""
- import json
-
- from osmosis_ai.rollout import RolloutRequest
- from osmosis_ai.rollout.core.base import RolloutContext
-
- debug_dir = str(tmp_path / "debug_logs")
-
- request = RolloutRequest(
- rollout_id="test-rollout-456",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
- mock_llm = MagicMock()
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- _debug_dir=debug_dir,
- )
-
- # Log some events
- ctx.log_event("pre_llm", turn=0, num_messages=3)
- ctx.log_event("llm_response", turn=0, has_tool_calls=True)
- ctx.log_event("rollout_complete", finish_reason="stop", reward=1.0)
-
- # Verify file was created
- import os
-
- debug_file = os.path.join(debug_dir, "test-rollout-456.jsonl")
- assert os.path.exists(debug_file)
-
- # Verify contents
- with open(debug_file, encoding="utf-8") as f:
- lines = f.readlines()
-
- assert len(lines) == 3
-
- event1 = json.loads(lines[0])
- assert event1["event"] == "pre_llm"
- assert event1["rollout_id"] == "test-rollout-456"
- assert event1["turn"] == 0
- assert event1["num_messages"] == 3
-
- event2 = json.loads(lines[1])
- assert event2["event"] == "llm_response"
- assert event2["has_tool_calls"] is True
-
- event3 = json.loads(lines[2])
- assert event3["event"] == "rollout_complete"
- assert event3["finish_reason"] == "stop"
- assert event3["reward"] == 1.0
-
-
-def test_rollout_context_get_debug_file_path() -> None:
- """Verify _get_debug_file_path returns correct path."""
- from osmosis_ai.rollout import RolloutRequest
- from osmosis_ai.rollout.core.base import RolloutContext
-
- request = RolloutRequest(
- rollout_id="my-rollout-id",
- server_url="http://localhost:8080",
- messages=[],
- completion_params={},
- )
-
- mock_llm = MagicMock()
-
- # Without debug_dir
- ctx_no_debug = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- )
- assert ctx_no_debug._get_debug_file_path() is None
-
- # With debug_dir
- ctx_with_debug = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- _debug_dir="/var/log/rollouts",
- )
- assert (
- ctx_with_debug._get_debug_file_path() == "/var/log/rollouts/my-rollout-id.jsonl"
- )
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-def test_create_app_with_debug_dir(tmp_path: Any) -> None:
- """Verify create_app accepts debug_dir parameter."""
- debug_dir = str(tmp_path / "debug_logs")
- agent_loop = SimpleAgentLoop()
- app = create_app(agent_loop, debug_dir=debug_dir)
-
- assert app is not None
-
-
-@pytest.mark.skipif(not FASTAPI_AVAILABLE, reason="FastAPI not installed")
-@pytest.mark.asyncio
-async def test_background_rollout_with_debug_logging(tmp_path: Any) -> None:
- """Verify debug logging works in background rollout execution."""
- import json
- import os
-
- from osmosis_ai.rollout import RolloutMetrics
-
- debug_dir = str(tmp_path / "debug_logs")
-
- class DebugLoggingAgentLoop(RolloutAgentLoop):
- """Agent loop that uses debug logging."""
-
- name = "debug_logging_agent"
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- # Log some events
- ctx.log_event("pre_llm", turn=0, num_messages=len(ctx.request.messages))
- ctx.log_event("rollout_complete", finish_reason="stop")
- return ctx.complete(list(ctx.request.messages))
-
- agent_loop = DebugLoggingAgentLoop()
- app = create_app(agent_loop, debug_dir=debug_dir)
-
- with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockClient:
- mock_client_instance = AsyncMock()
- mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
- mock_client_instance.__aexit__ = AsyncMock(return_value=None)
- mock_client_instance.complete_rollout = AsyncMock()
- mock_client_instance.get_metrics = MagicMock(return_value=RolloutMetrics())
- MockClient.return_value = mock_client_instance
-
- with TestClient(app) as client:
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "debug-test-789",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "test"}],
- "completion_params": {},
- },
- )
-
- assert response.status_code == 202
-
- # Give background task time to complete
- await asyncio.sleep(0.1)
-
- # Verify debug file was created
- debug_file = os.path.join(debug_dir, "debug-test-789.jsonl")
- assert os.path.exists(debug_file)
-
- # Verify contents
- with open(debug_file, encoding="utf-8") as f:
- lines = f.readlines()
-
- assert len(lines) == 2
-
- event1 = json.loads(lines[0])
- assert event1["event"] == "pre_llm"
- assert event1["rollout_id"] == "debug-test-789"
-
- event2 = json.loads(lines[1])
- assert event2["event"] == "rollout_complete"
diff --git a/tests/unit/test_litellm_provider.py b/tests/unit/test_litellm_provider.py
deleted file mode 100644
index 4a9d5075..00000000
--- a/tests/unit/test_litellm_provider.py
+++ /dev/null
@@ -1,490 +0,0 @@
-"""Tests for LiteLLM-based rubric evaluation."""
-
-import json
-import sys
-from unittest.mock import MagicMock, patch
-
-import litellm
-import pytest
-
-from osmosis_ai.rubric_eval import (
- DEFAULT_REQUEST_TIMEOUT_SECONDS,
- _default_timeout_for_model,
- _sanitize_json,
- _to_litellm_model,
- evaluate_rubric,
-)
-from osmosis_ai.rubric_types import ModelNotFoundError, ProviderRequestError
-
-
-class TestToLitellmModel:
- """Tests for model format conversion."""
-
- def test_openai_gpt5_uses_responses_prefix(self):
- assert (
- _to_litellm_model("openai", "gpt-5-mini") == "openai/responses/gpt-5-mini"
- )
- assert _to_litellm_model("openai", "gpt-5") == "openai/responses/gpt-5"
- assert (
- _to_litellm_model("OpenAI", "GPT-5-turbo") == "openai/responses/GPT-5-turbo"
- )
-
- def test_openai_other_models_standard_prefix(self):
- assert _to_litellm_model("openai", "gpt-4o") == "openai/gpt-4o"
- assert _to_litellm_model("openai", "gpt-4-turbo") == "openai/gpt-4-turbo"
- assert _to_litellm_model("openai", "o1-preview") == "openai/o1-preview"
-
- def test_anthropic_models(self):
- assert (
- _to_litellm_model("anthropic", "claude-sonnet-4-5-20250929")
- == "anthropic/claude-sonnet-4-5-20250929"
- )
- assert (
- _to_litellm_model("anthropic", "claude-3-opus-20240229")
- == "anthropic/claude-3-opus-20240229"
- )
-
- def test_xai_models(self):
- assert _to_litellm_model("xai", "grok-4-fast") == "xai/grok-4-fast"
- assert _to_litellm_model("xai", "grok-2") == "xai/grok-2"
-
- def test_gemini_models(self):
- assert (
- _to_litellm_model("gemini", "gemini-2.0-flash") == "gemini/gemini-2.0-flash"
- )
- assert _to_litellm_model("gemini", "gemini-1.5-pro") == "gemini/gemini-1.5-pro"
-
- def test_other_providers(self):
- assert (
- _to_litellm_model("cerebras", "llama-3.1-70b") == "cerebras/llama-3.1-70b"
- )
- assert (
- _to_litellm_model("openrouter", "meta-llama/llama-3-70b")
- == "openrouter/meta-llama/llama-3-70b"
- )
-
- def test_empty_provider_returns_model_only(self):
- assert _to_litellm_model("", "gpt-4o") == "gpt-4o"
-
- def test_case_insensitive(self):
- assert (
- _to_litellm_model("OPENAI", "gpt-5-mini") == "openai/responses/gpt-5-mini"
- )
- assert (
- _to_litellm_model("Anthropic", "claude-3-opus") == "anthropic/claude-3-opus"
- )
-
-
-class TestDefaultTimeoutForModel:
- """Tests for default timeout calculation."""
-
- def test_xai_grok4_timeout(self):
- assert _default_timeout_for_model("xai", "grok-4-fast") == 60.0
- assert _default_timeout_for_model("xai", "grok-4") == 60.0
-
- def test_xai_other_timeout(self):
- assert _default_timeout_for_model("xai", "grok-2") == 45.0
-
- def test_openai_gpt5_timeout(self):
- assert _default_timeout_for_model("openai", "gpt-5-mini") == 45.0
- assert _default_timeout_for_model("openai", "gpt-5") == 45.0
-
- def test_openai_other_timeout(self):
- assert (
- _default_timeout_for_model("openai", "gpt-4o")
- == DEFAULT_REQUEST_TIMEOUT_SECONDS
- )
-
- def test_gemini_timeout(self):
- assert _default_timeout_for_model("gemini", "gemini-2.0-flash") == 45.0
-
- def test_cerebras_timeout(self):
- assert _default_timeout_for_model("cerebras", "llama-4-scout-17b") == 60.0
-
- def test_openrouter_timeout(self):
- assert (
- _default_timeout_for_model("openrouter", "meta-llama/llama-4-maverick")
- == 60.0
- )
-
- def test_anthropic_timeout(self):
- assert (
- _default_timeout_for_model("anthropic", "claude-3-opus")
- == DEFAULT_REQUEST_TIMEOUT_SECONDS
- )
-
-
-class TestSanitizeJson:
- """Tests for JSON response parsing."""
-
- def test_valid_json(self):
- score, explanation = _sanitize_json(
- '{"score": 0.85, "explanation": "Good response"}'
- )
- assert score == 0.85
- assert explanation == "Good response"
-
- def test_json_with_code_fence(self):
- score, explanation = _sanitize_json(
- '```json\n{"score": 0.9, "explanation": "Test"}\n```'
- )
- assert score == 0.9
- assert explanation == "Test"
-
- def test_invalid_json_raises(self):
- with pytest.raises(ValueError, match="not valid JSON"):
- _sanitize_json("not valid json")
-
- def test_missing_score_raises(self):
- with pytest.raises(ValueError, match="numeric 'score'"):
- _sanitize_json('{"explanation": "No score"}')
-
- def test_missing_explanation_raises(self):
- with pytest.raises(ValueError, match="'explanation'"):
- _sanitize_json('{"score": 0.5}')
-
-
-def _create_mock_litellm_response(score: float, explanation: str) -> MagicMock:
- """Helper to create a mock LiteLLM response."""
- mock_response = MagicMock()
- mock_response.model_dump.return_value = {
- "choices": [
- {
- "message": {
- "content": json.dumps({"score": score, "explanation": explanation})
- }
- }
- ]
- }
- return mock_response
-
-
-class TestEvaluateRubric:
- """Tests for the main evaluate_rubric function."""
-
- def test_evaluate_rubric_success(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.85, "Good response"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score accuracy",
- solution_str="The answer is 42",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert result == 0.85
- mock_litellm.completion.assert_called_once()
- call_kwargs = mock_litellm.completion.call_args.kwargs
- assert call_kwargs["model"] == "openai/gpt-4o"
- assert call_kwargs["api_key"] == "test-key"
- assert call_kwargs["temperature"] == 0
-
- def test_evaluate_rubric_gpt5_no_temperature(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.9, "Excellent"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- evaluate_rubric(
- rubric="Score quality",
- solution_str="Test response",
- model_info={
- "provider": "openai",
- "model": "gpt-5-mini",
- "api_key": "test-key",
- },
- )
-
- call_kwargs = mock_litellm.completion.call_args.kwargs
- assert call_kwargs["model"] == "openai/responses/gpt-5-mini"
- assert "temperature" not in call_kwargs
-
- def test_evaluate_rubric_anthropic(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.9, "Excellent"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score quality",
- solution_str="Well written response",
- model_info={
- "provider": "anthropic",
- "model": "claude-sonnet-4-5-20250929",
- "api_key": "test-key",
- },
- )
-
- assert result == 0.9
- call_kwargs = mock_litellm.completion.call_args.kwargs
- assert call_kwargs["model"] == "anthropic/claude-sonnet-4-5-20250929"
-
- def test_evaluate_rubric_xai(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.8, "Good"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score response",
- solution_str="Some response",
- model_info={
- "provider": "xai",
- "model": "grok-4-fast",
- "api_key": "test-key",
- },
- )
-
- assert result == 0.8
- call_kwargs = mock_litellm.completion.call_args.kwargs
- assert call_kwargs["model"] == "xai/grok-4-fast"
-
- def test_evaluate_rubric_return_details(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.65, "Detailed explanation"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score accuracy",
- solution_str="Some answer",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- return_details=True,
- )
-
- assert result["score"] == 0.65
- assert result["explanation"] == "Detailed explanation"
- assert "raw" in result
-
- def test_evaluate_rubric_score_clamping_max(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 1.5, "Over max"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert result == 1.0 # Clamped to max
-
- def test_evaluate_rubric_score_clamping_min(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- -0.5, "Under min"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert result == 0.0 # Clamped to min
-
- def test_evaluate_rubric_api_error(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.APIError(
- message="API rate limit exceeded",
- llm_provider="openai",
- model="gpt-4o",
- status_code=429,
- )
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert "rate limit" in str(exc_info.value).lower()
-
- def test_evaluate_rubric_invalid_json_response(self):
- mock_litellm = MagicMock()
- mock_response = MagicMock()
- mock_response.model_dump.return_value = {
- "choices": [{"message": {"content": "Not valid JSON"}}]
- }
- mock_litellm.completion.return_value = mock_response
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert "not valid JSON" in str(exc_info.value)
-
- def test_evaluate_rubric_empty_response(self):
- mock_litellm = MagicMock()
- mock_response = MagicMock()
- mock_response.model_dump.return_value = {
- "choices": [{"message": {"content": ""}}]
- }
- mock_litellm.completion.return_value = mock_response
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert "did not include any content" in str(exc_info.value)
-
- def test_evaluate_rubric_model_not_found_404(self):
- """NotFoundError from litellm should raise ModelNotFoundError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.NotFoundError(
- message="The model `no-such-model` does not exist",
- llm_provider="openai",
- model="no-such-model",
- )
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ModelNotFoundError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "no-such-model",
- "api_key": "test-key",
- },
- )
-
- assert exc_info.value.provider == "openai"
- assert exc_info.value.model == "no-such-model"
-
- def test_evaluate_rubric_model_not_found_is_provider_request_error(self):
- """ModelNotFoundError should be catchable as ProviderRequestError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.NotFoundError(
- message="Not found",
- llm_provider="anthropic",
- model="bad-model",
- )
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError):
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "anthropic",
- "model": "bad-model",
- "api_key": "test-key",
- },
- )
-
- def test_evaluate_rubric_non_404_error_is_not_model_not_found(self):
- """Non-NotFoundError should raise ProviderRequestError, not ModelNotFoundError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.RateLimitError(
- message="Rate limit exceeded",
- llm_provider="openai",
- model="gpt-4o",
- )
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
- assert not isinstance(exc_info.value, ModelNotFoundError)
-
- def test_evaluate_rubric_custom_score_range(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 7.5, "Good"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score from 0-10",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- score_min=0.0,
- score_max=10.0,
- )
-
- assert result == 7.5
diff --git a/tests/unit/test_rubric_eval.py b/tests/unit/test_rubric_eval.py
index 7ecc3f4a..be03e66c 100644
--- a/tests/unit/test_rubric_eval.py
+++ b/tests/unit/test_rubric_eval.py
@@ -1,4 +1,26 @@
-from osmosis_ai.rubric_eval import _build_user_prompt, _select_text
+"""Tests for osmosis_ai.rubric_eval — prompt building, JSON parsing, and LiteLLM evaluation."""
+
+import json
+import sys
+from unittest.mock import MagicMock, patch
+
+import litellm
+import pytest
+
+from osmosis_ai.rubric_eval import (
+ DEFAULT_REQUEST_TIMEOUT_SECONDS,
+ _build_user_prompt,
+ _default_timeout_for_model,
+ _sanitize_json,
+ _select_text,
+ _to_litellm_model,
+ evaluate_rubric,
+)
+from osmosis_ai.rubric_types import ModelNotFoundError, ProviderRequestError
+
+# =============================================================================
+# _select_text / _build_user_prompt Tests
+# =============================================================================
def test_build_user_prompt_basic_blocks() -> None:
@@ -38,3 +60,496 @@ def test_build_user_prompt_with_optional_sections() -> None:
def test_select_text_prefers_first_non_empty() -> None:
assert _select_text(None, "", " value ", "fallback") == "value"
assert _select_text(None, " ") is None
+
+
+# =============================================================================
+# _to_litellm_model Tests
+# =============================================================================
+
+
+class TestToLitellmModel:
+ """Tests for model format conversion."""
+
+ def test_openai_gpt5_uses_responses_prefix(self):
+ assert (
+ _to_litellm_model("openai", "gpt-5-mini") == "openai/responses/gpt-5-mini"
+ )
+ assert _to_litellm_model("openai", "gpt-5") == "openai/responses/gpt-5"
+ assert (
+ _to_litellm_model("OpenAI", "GPT-5-turbo") == "openai/responses/GPT-5-turbo"
+ )
+
+ def test_openai_other_models_standard_prefix(self):
+ assert _to_litellm_model("openai", "gpt-4o") == "openai/gpt-4o"
+ assert _to_litellm_model("openai", "gpt-4-turbo") == "openai/gpt-4-turbo"
+ assert _to_litellm_model("openai", "o1-preview") == "openai/o1-preview"
+
+ def test_anthropic_models(self):
+ assert (
+ _to_litellm_model("anthropic", "claude-sonnet-4-5-20250929")
+ == "anthropic/claude-sonnet-4-5-20250929"
+ )
+ assert (
+ _to_litellm_model("anthropic", "claude-3-opus-20240229")
+ == "anthropic/claude-3-opus-20240229"
+ )
+
+ def test_xai_models(self):
+ assert _to_litellm_model("xai", "grok-4-fast") == "xai/grok-4-fast"
+ assert _to_litellm_model("xai", "grok-2") == "xai/grok-2"
+
+ def test_gemini_models(self):
+ assert (
+ _to_litellm_model("gemini", "gemini-2.0-flash") == "gemini/gemini-2.0-flash"
+ )
+ assert _to_litellm_model("gemini", "gemini-1.5-pro") == "gemini/gemini-1.5-pro"
+
+ def test_other_providers(self):
+ assert (
+ _to_litellm_model("cerebras", "llama-3.1-70b") == "cerebras/llama-3.1-70b"
+ )
+ assert (
+ _to_litellm_model("openrouter", "meta-llama/llama-3-70b")
+ == "openrouter/meta-llama/llama-3-70b"
+ )
+
+ def test_empty_provider_returns_model_only(self):
+ assert _to_litellm_model("", "gpt-4o") == "gpt-4o"
+
+ def test_case_insensitive(self):
+ assert (
+ _to_litellm_model("OPENAI", "gpt-5-mini") == "openai/responses/gpt-5-mini"
+ )
+ assert (
+ _to_litellm_model("Anthropic", "claude-3-opus") == "anthropic/claude-3-opus"
+ )
+
+
+# =============================================================================
+# _default_timeout_for_model Tests
+# =============================================================================
+
+
+class TestDefaultTimeoutForModel:
+ """Tests for default timeout calculation."""
+
+ def test_xai_grok4_timeout(self):
+ assert _default_timeout_for_model("xai", "grok-4-fast") == 60.0
+ assert _default_timeout_for_model("xai", "grok-4") == 60.0
+
+ def test_xai_other_timeout(self):
+ assert _default_timeout_for_model("xai", "grok-2") == 45.0
+
+ def test_openai_gpt5_timeout(self):
+ assert _default_timeout_for_model("openai", "gpt-5-mini") == 45.0
+ assert _default_timeout_for_model("openai", "gpt-5") == 45.0
+
+ def test_openai_other_timeout(self):
+ assert (
+ _default_timeout_for_model("openai", "gpt-4o")
+ == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ )
+
+ def test_gemini_timeout(self):
+ assert _default_timeout_for_model("gemini", "gemini-2.0-flash") == 45.0
+
+ def test_cerebras_timeout(self):
+ assert _default_timeout_for_model("cerebras", "llama-4-scout-17b") == 60.0
+
+ def test_openrouter_timeout(self):
+ assert (
+ _default_timeout_for_model("openrouter", "meta-llama/llama-4-maverick")
+ == 60.0
+ )
+
+ def test_anthropic_timeout(self):
+ assert (
+ _default_timeout_for_model("anthropic", "claude-3-opus")
+ == DEFAULT_REQUEST_TIMEOUT_SECONDS
+ )
+
+
+# =============================================================================
+# _sanitize_json Tests
+# =============================================================================
+
+
+class TestSanitizeJson:
+ """Tests for JSON response parsing."""
+
+ def test_valid_json(self):
+ score, explanation = _sanitize_json(
+ '{"score": 0.85, "explanation": "Good response"}'
+ )
+ assert score == 0.85
+ assert explanation == "Good response"
+
+ def test_json_with_code_fence(self):
+ score, explanation = _sanitize_json(
+ '```json\n{"score": 0.9, "explanation": "Test"}\n```'
+ )
+ assert score == 0.9
+ assert explanation == "Test"
+
+ def test_invalid_json_raises(self):
+ with pytest.raises(ValueError, match="not valid JSON"):
+ _sanitize_json("not valid json")
+
+ def test_missing_score_raises(self):
+ with pytest.raises(ValueError, match="numeric 'score'"):
+ _sanitize_json('{"explanation": "No score"}')
+
+ def test_missing_explanation_raises(self):
+ with pytest.raises(ValueError, match="'explanation'"):
+ _sanitize_json('{"score": 0.5}')
+
+
+# =============================================================================
+# evaluate_rubric Tests
+# =============================================================================
+
+
+def _create_mock_litellm_response(score: float, explanation: str) -> MagicMock:
+ """Helper to create a mock LiteLLM response."""
+ mock_response = MagicMock()
+ mock_response.model_dump.return_value = {
+ "choices": [
+ {
+ "message": {
+ "content": json.dumps({"score": score, "explanation": explanation})
+ }
+ }
+ ]
+ }
+ return mock_response
+
+
+class TestEvaluateRubric:
+ """Tests for the main evaluate_rubric function."""
+
+ def test_evaluate_rubric_success(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.85, "Good response"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score accuracy",
+ solution_str="The answer is 42",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert result == 0.85
+ mock_litellm.completion.assert_called_once()
+ call_kwargs = mock_litellm.completion.call_args.kwargs
+ assert call_kwargs["model"] == "openai/gpt-4o"
+ assert call_kwargs["api_key"] == "test-key"
+ assert call_kwargs["temperature"] == 0
+
+ def test_evaluate_rubric_gpt5_no_temperature(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.9, "Excellent"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ evaluate_rubric(
+ rubric="Score quality",
+ solution_str="Test response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5-mini",
+ "api_key": "test-key",
+ },
+ )
+
+ call_kwargs = mock_litellm.completion.call_args.kwargs
+ assert call_kwargs["model"] == "openai/responses/gpt-5-mini"
+ assert "temperature" not in call_kwargs
+
+ def test_evaluate_rubric_anthropic(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.9, "Excellent"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score quality",
+ solution_str="Well written response",
+ model_info={
+ "provider": "anthropic",
+ "model": "claude-sonnet-4-5-20250929",
+ "api_key": "test-key",
+ },
+ )
+
+ assert result == 0.9
+ call_kwargs = mock_litellm.completion.call_args.kwargs
+ assert call_kwargs["model"] == "anthropic/claude-sonnet-4-5-20250929"
+
+ def test_evaluate_rubric_xai(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.8, "Good"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score response",
+ solution_str="Some response",
+ model_info={
+ "provider": "xai",
+ "model": "grok-4-fast",
+ "api_key": "test-key",
+ },
+ )
+
+ assert result == 0.8
+ call_kwargs = mock_litellm.completion.call_args.kwargs
+ assert call_kwargs["model"] == "xai/grok-4-fast"
+
+ def test_evaluate_rubric_return_details(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 0.65, "Detailed explanation"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score accuracy",
+ solution_str="Some answer",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ return_details=True,
+ )
+
+ assert result["score"] == 0.65
+ assert result["explanation"] == "Detailed explanation"
+ assert "raw" in result
+
+ def test_evaluate_rubric_score_clamping_max(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 1.5, "Over max"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert result == 1.0 # Clamped to max
+
+ def test_evaluate_rubric_score_clamping_min(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ -0.5, "Under min"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert result == 0.0 # Clamped to min
+
+ def test_evaluate_rubric_api_error(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.side_effect = litellm.APIError(
+ message="API rate limit exceeded",
+ llm_provider="openai",
+ model="gpt-4o",
+ status_code=429,
+ )
+ mock_litellm.APIError = litellm.APIError
+ mock_litellm.NotFoundError = litellm.NotFoundError
+ mock_litellm.RateLimitError = litellm.RateLimitError
+ mock_litellm.AuthenticationError = litellm.AuthenticationError
+ mock_litellm.Timeout = litellm.Timeout
+ mock_litellm.APIConnectionError = litellm.APIConnectionError
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert "rate limit" in str(exc_info.value).lower()
+
+ def test_evaluate_rubric_invalid_json_response(self):
+ mock_litellm = MagicMock()
+ mock_response = MagicMock()
+ mock_response.model_dump.return_value = {
+ "choices": [{"message": {"content": "Not valid JSON"}}]
+ }
+ mock_litellm.completion.return_value = mock_response
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert "not valid JSON" in str(exc_info.value)
+
+ def test_evaluate_rubric_empty_response(self):
+ mock_litellm = MagicMock()
+ mock_response = MagicMock()
+ mock_response.model_dump.return_value = {
+ "choices": [{"message": {"content": ""}}]
+ }
+ mock_litellm.completion.return_value = mock_response
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert "did not include any content" in str(exc_info.value)
+
+ def test_evaluate_rubric_model_not_found_404(self):
+ """NotFoundError from litellm should raise ModelNotFoundError."""
+ mock_litellm = MagicMock()
+ mock_litellm.completion.side_effect = litellm.NotFoundError(
+ message="The model `no-such-model` does not exist",
+ llm_provider="openai",
+ model="no-such-model",
+ )
+ mock_litellm.APIError = litellm.APIError
+ mock_litellm.NotFoundError = litellm.NotFoundError
+ mock_litellm.RateLimitError = litellm.RateLimitError
+ mock_litellm.AuthenticationError = litellm.AuthenticationError
+ mock_litellm.Timeout = litellm.Timeout
+ mock_litellm.APIConnectionError = litellm.APIConnectionError
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ModelNotFoundError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "no-such-model",
+ "api_key": "test-key",
+ },
+ )
+
+ assert exc_info.value.provider == "openai"
+ assert exc_info.value.model == "no-such-model"
+
+ def test_evaluate_rubric_model_not_found_is_provider_request_error(self):
+ """ModelNotFoundError should be catchable as ProviderRequestError."""
+ mock_litellm = MagicMock()
+ mock_litellm.completion.side_effect = litellm.NotFoundError(
+ message="Not found",
+ llm_provider="anthropic",
+ model="bad-model",
+ )
+ mock_litellm.APIError = litellm.APIError
+ mock_litellm.NotFoundError = litellm.NotFoundError
+ mock_litellm.RateLimitError = litellm.RateLimitError
+ mock_litellm.AuthenticationError = litellm.AuthenticationError
+ mock_litellm.Timeout = litellm.Timeout
+ mock_litellm.APIConnectionError = litellm.APIConnectionError
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ProviderRequestError):
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "anthropic",
+ "model": "bad-model",
+ "api_key": "test-key",
+ },
+ )
+
+ def test_evaluate_rubric_non_404_error_is_not_model_not_found(self):
+ """Non-NotFoundError should raise ProviderRequestError, not ModelNotFoundError."""
+ mock_litellm = MagicMock()
+ mock_litellm.completion.side_effect = litellm.RateLimitError(
+ message="Rate limit exceeded",
+ llm_provider="openai",
+ model="gpt-4o",
+ )
+ mock_litellm.APIError = litellm.APIError
+ mock_litellm.NotFoundError = litellm.NotFoundError
+ mock_litellm.RateLimitError = litellm.RateLimitError
+ mock_litellm.AuthenticationError = litellm.AuthenticationError
+ mock_litellm.Timeout = litellm.Timeout
+ mock_litellm.APIConnectionError = litellm.APIConnectionError
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
+
+ assert not isinstance(exc_info.value, ModelNotFoundError)
+
+ def test_evaluate_rubric_custom_score_range(self):
+ mock_litellm = MagicMock()
+ mock_litellm.completion.return_value = _create_mock_litellm_response(
+ 7.5, "Good"
+ )
+
+ with patch.dict(sys.modules, {"litellm": mock_litellm}):
+ result = evaluate_rubric(
+ rubric="Score from 0-10",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ score_min=0.0,
+ score_max=10.0,
+ )
+
+ assert result == 7.5
diff --git a/tests/unit/test_utils_decorators.py b/tests/unit/test_utils.py
similarity index 100%
rename from tests/unit/test_utils_decorators.py
rename to tests/unit/test_utils.py
From d13be797578cf499b6d6f8e31fafe3079ebce977 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 14:45:45 -0800
Subject: [PATCH 17/60] bugfix
---
tests/__init__.py | 0
1 file changed, 0 insertions(+), 0 deletions(-)
create mode 100644 tests/__init__.py
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
From 5e427efd125cd66f97dc476a42a2d14291d39430 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 15:19:42 -0800
Subject: [PATCH 18/60] Enhance pytest configuration and refactor test fixtures
- Added custom markers for integration and slow tests in `pytest.ini`.
- Refactored test fixtures in `conftest.py` to load sample data from JSON files, improving maintainability and clarity.
- Removed unused mock agent loop classes to streamline the test setup.
- Updated tests to utilize the new fixture structure, ensuring consistency and better organization.
---
pytest.ini | 7 +-
tests/_data/__init__.py | 11 ++
tests/_data/sample_messages.json | 29 +++
tests/_data/sample_rollout_request.json | 11 ++
tests/_data/sample_tool_schema.json | 15 ++
tests/_helpers/__init__.py | 18 ++
tests/_helpers/factories.py | 36 ++++
tests/_helpers/mock_agents.py | 82 +++++++++
tests/conftest.py | 170 ++----------------
tests/integration/__init__.py | 0
tests/integration/cli/__init__.py | 0
tests/integration/cli/test_cli_integration.py | 23 +++
tests/integration/conftest.py | 15 ++
tests/integration/rollout/__init__.py | 0
.../integration/rollout/test_rollout_flow.py | 23 +++
tests/integration/server/__init__.py | 0
.../server/test_app_integration.py | 23 +++
tests/unit/rollout/server/test_app.py | 4 +-
18 files changed, 304 insertions(+), 163 deletions(-)
create mode 100644 tests/_data/__init__.py
create mode 100644 tests/_data/sample_messages.json
create mode 100644 tests/_data/sample_rollout_request.json
create mode 100644 tests/_data/sample_tool_schema.json
create mode 100644 tests/_helpers/__init__.py
create mode 100644 tests/_helpers/factories.py
create mode 100644 tests/_helpers/mock_agents.py
create mode 100644 tests/integration/__init__.py
create mode 100644 tests/integration/cli/__init__.py
create mode 100644 tests/integration/cli/test_cli_integration.py
create mode 100644 tests/integration/conftest.py
create mode 100644 tests/integration/rollout/__init__.py
create mode 100644 tests/integration/rollout/test_rollout_flow.py
create mode 100644 tests/integration/server/__init__.py
create mode 100644 tests/integration/server/test_app_integration.py
diff --git a/pytest.ini b/pytest.ini
index 04d3b842..88a197d9 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,10 +1,11 @@
[pytest]
-# Suppress all PytestAssertRewriteWarning warnings
filterwarnings =
ignore::pytest.PytestAssertRewriteWarning
-# Configure asyncio mode for pytest-asyncio
asyncio_mode = auto
-# Explicitly define test discovery path
testpaths = tests
+
+markers =
+ integration: tests requiring real component interaction (deselect with -m "not integration")
+ slow: tests that take more than 5 seconds
diff --git a/tests/_data/__init__.py b/tests/_data/__init__.py
new file mode 100644
index 00000000..6c81c6a2
--- /dev/null
+++ b/tests/_data/__init__.py
@@ -0,0 +1,11 @@
+"""Test fixture data directory.
+
+Usage:
+ from tests._data import DATA_DIR
+
+ path = DATA_DIR / "sample_rollout_request.json"
+"""
+
+from pathlib import Path
+
+DATA_DIR = Path(__file__).parent
diff --git a/tests/_data/sample_messages.json b/tests/_data/sample_messages.json
new file mode 100644
index 00000000..32e92676
--- /dev/null
+++ b/tests/_data/sample_messages.json
@@ -0,0 +1,29 @@
+{
+ "system_user_pair": [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "What is 2 + 2?"}
+ ],
+ "assistant_response": {
+ "role": "assistant",
+ "content": "The answer is 4."
+ },
+ "assistant_with_tool_calls": {
+ "role": "assistant",
+ "content": null,
+ "tool_calls": [
+ {
+ "id": "call_123",
+ "type": "function",
+ "function": {
+ "name": "add",
+ "arguments": "{\"a\": 2, \"b\": 2}"
+ }
+ }
+ ]
+ },
+ "tool_result": {
+ "role": "tool",
+ "content": "4",
+ "tool_call_id": "call_123"
+ }
+}
diff --git a/tests/_data/sample_rollout_request.json b/tests/_data/sample_rollout_request.json
new file mode 100644
index 00000000..49426a04
--- /dev/null
+++ b/tests/_data/sample_rollout_request.json
@@ -0,0 +1,11 @@
+{
+ "rollout_id": "test-rollout-123",
+ "server_url": "http://localhost:8080",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ],
+ "completion_params": {
+ "temperature": 0.7,
+ "max_tokens": 512
+ }
+}
diff --git a/tests/_data/sample_tool_schema.json b/tests/_data/sample_tool_schema.json
new file mode 100644
index 00000000..75acb507
--- /dev/null
+++ b/tests/_data/sample_tool_schema.json
@@ -0,0 +1,15 @@
+{
+ "type": "function",
+ "function": {
+ "name": "add",
+ "description": "Add two numbers",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "a": {"type": "number"},
+ "b": {"type": "number"}
+ },
+ "required": ["a", "b"]
+ }
+ }
+}
diff --git a/tests/_helpers/__init__.py b/tests/_helpers/__init__.py
new file mode 100644
index 00000000..51d4cae1
--- /dev/null
+++ b/tests/_helpers/__init__.py
@@ -0,0 +1,18 @@
+"""Shared test helpers — mock agents, factory functions, and utilities."""
+
+from tests._helpers.factories import make_rollout_payload, mock_llm_client
+from tests._helpers.mock_agents import (
+ FailingAgentLoop,
+ MockAgentLoop,
+ SimpleAgentLoop,
+ SlowAgentLoop,
+)
+
+__all__ = [
+ "FailingAgentLoop",
+ "MockAgentLoop",
+ "SimpleAgentLoop",
+ "SlowAgentLoop",
+ "make_rollout_payload",
+ "mock_llm_client",
+]
diff --git a/tests/_helpers/factories.py b/tests/_helpers/factories.py
new file mode 100644
index 00000000..97bb50f6
--- /dev/null
+++ b/tests/_helpers/factories.py
@@ -0,0 +1,36 @@
+"""Reusable factory / helper functions for tests."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+from osmosis_ai.rollout import RolloutMetrics
+
+
+def make_rollout_payload(
+ rollout_id: str = "test-123",
+ idempotency_key: str | None = None,
+ **overrides: Any,
+) -> dict[str, Any]:
+ """Build a minimal valid RolloutRequest JSON payload."""
+ payload: dict[str, Any] = {
+ "rollout_id": rollout_id,
+ "server_url": "http://localhost:8080",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "completion_params": {"temperature": 0.7},
+ }
+ if idempotency_key is not None:
+ payload["idempotency_key"] = idempotency_key
+ payload.update(overrides)
+ return payload
+
+
+def mock_llm_client() -> MagicMock:
+ """Create a mocked OsmosisLLMClient usable as an async context manager."""
+ mock = AsyncMock()
+ mock.__aenter__ = AsyncMock(return_value=mock)
+ mock.__aexit__ = AsyncMock(return_value=None)
+ mock.complete_rollout = AsyncMock()
+ mock.get_metrics = MagicMock(return_value=RolloutMetrics())
+ return mock
diff --git a/tests/_helpers/mock_agents.py b/tests/_helpers/mock_agents.py
new file mode 100644
index 00000000..9b6f6196
--- /dev/null
+++ b/tests/_helpers/mock_agents.py
@@ -0,0 +1,82 @@
+"""Reusable mock agent loop implementations for tests."""
+
+from __future__ import annotations
+
+import asyncio
+
+from osmosis_ai.rollout import (
+ OpenAIFunctionToolSchema,
+ RolloutAgentLoop,
+ RolloutContext,
+ RolloutRequest,
+ RolloutResult,
+)
+
+
+class MockAgentLoop(RolloutAgentLoop):
+ """Mock agent loop for testing."""
+
+ name = "mock_agent"
+
+ def __init__(
+ self,
+ tools: list[OpenAIFunctionToolSchema] | None = None,
+ run_result: RolloutResult | None = None,
+ run_error: Exception | None = None,
+ ):
+ self._tools = tools or []
+ self._run_result = run_result
+ self._run_error = run_error
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return self._tools
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ if self._run_error:
+ raise self._run_error
+ if self._run_result:
+ return self._run_result
+ return ctx.complete(list(ctx.request.messages))
+
+
+class SimpleAgentLoop(RolloutAgentLoop):
+ """Simple agent loop that completes immediately."""
+
+ name = "simple_agent"
+
+ def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
+ self._tools = tools or []
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return self._tools
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ return ctx.complete(list(ctx.request.messages))
+
+
+class FailingAgentLoop(RolloutAgentLoop):
+ """Agent loop that raises an error."""
+
+ name = "failing_agent"
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ raise RuntimeError("Test error")
+
+
+class SlowAgentLoop(RolloutAgentLoop):
+ """Agent loop that sleeps for a configurable duration before completing."""
+
+ name = "slow_agent"
+
+ def __init__(self, delay: float = 0.5):
+ self._delay = delay
+
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ await asyncio.sleep(self._delay)
+ return ctx.complete(list(ctx.request.messages))
diff --git a/tests/conftest.py b/tests/conftest.py
index 10699cac..41ff4c72 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,49 +2,21 @@
from __future__ import annotations
-import asyncio
+import json
from typing import Any
-from unittest.mock import AsyncMock, MagicMock
import pytest
from osmosis_ai.rollout import (
- OpenAIFunctionParametersSchema,
- OpenAIFunctionPropertySchema,
- OpenAIFunctionSchema,
OpenAIFunctionToolSchema,
- RolloutAgentLoop,
- RolloutContext,
- RolloutMetrics,
RolloutRequest,
- RolloutResult,
)
+from tests._data import DATA_DIR
+from tests._helpers import MockAgentLoop
-class MockAgentLoop(RolloutAgentLoop):
- """Mock agent loop for testing."""
-
- name = "mock_agent"
-
- def __init__(
- self,
- tools: list[OpenAIFunctionToolSchema] | None = None,
- run_result: RolloutResult | None = None,
- run_error: Exception | None = None,
- ):
- self._tools = tools or []
- self._run_result = run_result
- self._run_error = run_error
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return self._tools
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- if self._run_error:
- raise self._run_error
- if self._run_result:
- return self._run_result
- return ctx.complete(list(ctx.request.messages))
+def _load_json(name: str) -> Any:
+ return json.loads((DATA_DIR / name).read_text())
@pytest.fixture
@@ -56,154 +28,34 @@ def mock_agent_loop() -> MockAgentLoop:
@pytest.fixture
def sample_rollout_request() -> RolloutRequest:
"""Create a sample RolloutRequest for testing."""
- return RolloutRequest(
- rollout_id="test-rollout-123",
- server_url="http://localhost:8080",
- messages=[{"role": "user", "content": "Hello"}],
- completion_params={"temperature": 0.7, "max_tokens": 512},
- )
+ return RolloutRequest(**_load_json("sample_rollout_request.json"))
@pytest.fixture
def sample_tool_schema() -> OpenAIFunctionToolSchema:
"""Create a sample OpenAIFunctionToolSchema for testing."""
- return OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="add",
- description="Add two numbers",
- parameters=OpenAIFunctionParametersSchema(
- type="object",
- properties={
- "a": OpenAIFunctionPropertySchema(type="number"),
- "b": OpenAIFunctionPropertySchema(type="number"),
- },
- required=["a", "b"],
- ),
- ),
- )
+ return OpenAIFunctionToolSchema(**_load_json("sample_tool_schema.json"))
@pytest.fixture
def sample_messages() -> list[dict[str, Any]]:
"""Create sample messages for testing."""
- return [
- {"role": "system", "content": "You are a helpful assistant."},
- {"role": "user", "content": "What is 2 + 2?"},
- ]
+ return _load_json("sample_messages.json")["system_user_pair"]
@pytest.fixture
def sample_assistant_message() -> dict[str, Any]:
"""Create a sample assistant message for testing."""
- return {"role": "assistant", "content": "The answer is 4."}
+ return _load_json("sample_messages.json")["assistant_response"]
@pytest.fixture
def sample_assistant_message_with_tool_calls() -> dict[str, Any]:
"""Create a sample assistant message with tool calls."""
- return {
- "role": "assistant",
- "content": None,
- "tool_calls": [
- {
- "id": "call_123",
- "type": "function",
- "function": {
- "name": "add",
- "arguments": '{"a": 2, "b": 2}',
- },
- }
- ],
- }
+ return _load_json("sample_messages.json")["assistant_with_tool_calls"]
@pytest.fixture
def sample_tool_message() -> dict[str, Any]:
"""Create a sample tool result message."""
- return {
- "role": "tool",
- "content": "4",
- "tool_call_id": "call_123",
- }
-
-
-# =============================================================================
-# Shared agent loop stubs (used by multiple test modules)
-# =============================================================================
-
-
-class SimpleAgentLoop(RolloutAgentLoop):
- """Simple agent loop that completes immediately."""
-
- name = "simple_agent"
-
- def __init__(self, tools: list[OpenAIFunctionToolSchema] | None = None):
- self._tools = tools or []
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return self._tools
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- return ctx.complete(list(ctx.request.messages))
-
-
-class FailingAgentLoop(RolloutAgentLoop):
- """Agent loop that raises an error."""
-
- name = "failing_agent"
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- raise RuntimeError("Test error")
-
-
-class SlowAgentLoop(RolloutAgentLoop):
- """Agent loop that sleeps for a configurable duration before completing."""
-
- name = "slow_agent"
-
- def __init__(self, delay: float = 0.5):
- self._delay = delay
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- await asyncio.sleep(self._delay)
- return ctx.complete(list(ctx.request.messages))
-
-
-# =============================================================================
-# Shared helper functions (used by multiple test modules)
-# =============================================================================
-
-
-def make_rollout_payload(
- rollout_id: str = "test-123",
- idempotency_key: str | None = None,
- **overrides: Any,
-) -> dict[str, Any]:
- """Build a minimal valid RolloutRequest JSON payload."""
- payload: dict[str, Any] = {
- "rollout_id": rollout_id,
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- }
- if idempotency_key is not None:
- payload["idempotency_key"] = idempotency_key
- payload.update(overrides)
- return payload
-
-
-def mock_llm_client() -> MagicMock:
- """Create a mocked OsmosisLLMClient usable as an async context manager."""
- mock = AsyncMock()
- mock.__aenter__ = AsyncMock(return_value=mock)
- mock.__aexit__ = AsyncMock(return_value=None)
- mock.complete_rollout = AsyncMock()
- mock.get_metrics = MagicMock(return_value=RolloutMetrics())
- return mock
+ return _load_json("sample_messages.json")["tool_result"]
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/integration/cli/__init__.py b/tests/integration/cli/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/integration/cli/test_cli_integration.py b/tests/integration/cli/test_cli_integration.py
new file mode 100644
index 00000000..78d1ef99
--- /dev/null
+++ b/tests/integration/cli/test_cli_integration.py
@@ -0,0 +1,23 @@
+"""Integration tests for CLI commands.
+
+Tests CLI commands using click.testing.CliRunner with real file system
+interactions and minimal mocking.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+class TestCLICommands:
+ """Test CLI commands with real invocation."""
+
+ @pytest.mark.integration
+ def test_whoami_unauthenticated(self):
+ """osmosis whoami shows appropriate message when not logged in."""
+ pytest.skip("TODO: implement whoami integration test")
+
+ @pytest.mark.integration
+ def test_validate_command(self):
+ """osmosis validate checks agent module correctly."""
+ pytest.skip("TODO: implement validate integration test")
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
new file mode 100644
index 00000000..93610482
--- /dev/null
+++ b/tests/integration/conftest.py
@@ -0,0 +1,15 @@
+"""Fixtures specific to integration tests.
+
+Integration tests exercise real component interactions with minimal mocking.
+They may use httpx.AsyncClient against real FastAPI apps, real file I/O, etc.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+@pytest.fixture
+def anyio_backend():
+ """Use asyncio as the async backend for integration tests."""
+ return "asyncio"
diff --git a/tests/integration/rollout/__init__.py b/tests/integration/rollout/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/integration/rollout/test_rollout_flow.py b/tests/integration/rollout/test_rollout_flow.py
new file mode 100644
index 00000000..721cf0dc
--- /dev/null
+++ b/tests/integration/rollout/test_rollout_flow.py
@@ -0,0 +1,23 @@
+"""Integration tests for the complete rollout lifecycle.
+
+Tests the full init -> chat -> completed flow using real FastAPI app
+with httpx.AsyncClient, minimal mocking (only the LLM client).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+class TestRolloutLifecycle:
+ """Test the complete rollout request lifecycle."""
+
+ @pytest.mark.integration
+ async def test_init_chat_complete_flow(self):
+ """Full rollout: init request -> agent runs -> completion reported."""
+ pytest.skip("TODO: implement rollout lifecycle integration test")
+
+ @pytest.mark.integration
+ async def test_concurrent_rollouts(self):
+ """Multiple rollouts can execute concurrently without interference."""
+ pytest.skip("TODO: implement concurrent rollout integration test")
diff --git a/tests/integration/server/__init__.py b/tests/integration/server/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/integration/server/test_app_integration.py b/tests/integration/server/test_app_integration.py
new file mode 100644
index 00000000..2c40b6ab
--- /dev/null
+++ b/tests/integration/server/test_app_integration.py
@@ -0,0 +1,23 @@
+"""Integration tests for the FastAPI rollout server.
+
+Tests HTTP endpoints using httpx.AsyncClient against a real create_app() instance.
+Only the external LLM client is mocked.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+class TestServerEndpoints:
+ """Test server HTTP endpoints with real request/response cycle."""
+
+ @pytest.mark.integration
+ async def test_health_endpoint(self):
+ """GET /health returns 200 with server status."""
+ pytest.skip("TODO: implement health endpoint integration test")
+
+ @pytest.mark.integration
+ async def test_rollout_endpoint_validation(self):
+ """POST /v1/rollout/init validates request schema."""
+ pytest.skip("TODO: implement rollout validation integration test")
diff --git a/tests/unit/rollout/server/test_app.py b/tests/unit/rollout/server/test_app.py
index 0c5bed87..909aa95f 100644
--- a/tests/unit/rollout/server/test_app.py
+++ b/tests/unit/rollout/server/test_app.py
@@ -27,7 +27,7 @@
create_app,
)
from osmosis_ai.rollout.server.app import _extract_bearer_token
-from tests.conftest import (
+from tests._helpers import (
SimpleAgentLoop,
SlowAgentLoop,
make_rollout_payload,
@@ -230,6 +230,7 @@ def test_registration_not_started_when_port_is_none(self) -> None:
resp = client.get("/health")
assert resp.status_code == 200
+ @pytest.mark.slow
async def test_registration_task_started_with_credentials(self) -> None:
"""Registration background task should be created when credentials are given."""
mock_creds = MagicMock()
@@ -258,6 +259,7 @@ async def test_registration_task_started_with_credentials(self) -> None:
# Give the registration background task time to run
await asyncio.sleep(0.5)
+ @pytest.mark.slow
async def test_registration_error_handled_gracefully(self) -> None:
"""If the registration raises an error, shutdown should not crash."""
mock_creds = MagicMock()
From 345b1c3d0d1239088a1a1b9ad52ae37b42ce5299 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 15:33:07 -0800
Subject: [PATCH 19/60] Update GitHub Actions workflow to support multiple
Python versions for testing
- Added a strategy matrix to the `pytest` job in `tests.yml` to run tests on Python 3.10 and 3.12.
- Modified the Python setup step to dynamically use the specified Python version from the matrix.
---
.github/workflows/tests.yml | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 096db083..7edcd33f 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -30,15 +30,18 @@ jobs:
pytest:
runs-on: ubuntu-latest
+ strategy:
+ matrix:
+ python-version: ["3.10", "3.12"]
steps:
- name: Checkout repository
uses: actions/checkout@v4
- - name: Set up Python
+ - name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
- python-version: "3.12"
+ python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: pip install -e ".[dev]"
From ed1ba2480c55428b353beee2f5b7aec7abcaf794 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 16:05:50 -0800
Subject: [PATCH 20/60] Remove integration test files and related fixtures from
the test suite
---
tests/integration/__init__.py | 0
tests/integration/cli/__init__.py | 0
tests/integration/cli/test_cli_integration.py | 23 -------------------
tests/integration/conftest.py | 15 ------------
tests/integration/rollout/__init__.py | 0
.../integration/rollout/test_rollout_flow.py | 23 -------------------
tests/integration/server/__init__.py | 0
.../server/test_app_integration.py | 23 -------------------
8 files changed, 84 deletions(-)
delete mode 100644 tests/integration/__init__.py
delete mode 100644 tests/integration/cli/__init__.py
delete mode 100644 tests/integration/cli/test_cli_integration.py
delete mode 100644 tests/integration/conftest.py
delete mode 100644 tests/integration/rollout/__init__.py
delete mode 100644 tests/integration/rollout/test_rollout_flow.py
delete mode 100644 tests/integration/server/__init__.py
delete mode 100644 tests/integration/server/test_app_integration.py
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/tests/integration/cli/__init__.py b/tests/integration/cli/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/tests/integration/cli/test_cli_integration.py b/tests/integration/cli/test_cli_integration.py
deleted file mode 100644
index 78d1ef99..00000000
--- a/tests/integration/cli/test_cli_integration.py
+++ /dev/null
@@ -1,23 +0,0 @@
-"""Integration tests for CLI commands.
-
-Tests CLI commands using click.testing.CliRunner with real file system
-interactions and minimal mocking.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-
-class TestCLICommands:
- """Test CLI commands with real invocation."""
-
- @pytest.mark.integration
- def test_whoami_unauthenticated(self):
- """osmosis whoami shows appropriate message when not logged in."""
- pytest.skip("TODO: implement whoami integration test")
-
- @pytest.mark.integration
- def test_validate_command(self):
- """osmosis validate checks agent module correctly."""
- pytest.skip("TODO: implement validate integration test")
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
deleted file mode 100644
index 93610482..00000000
--- a/tests/integration/conftest.py
+++ /dev/null
@@ -1,15 +0,0 @@
-"""Fixtures specific to integration tests.
-
-Integration tests exercise real component interactions with minimal mocking.
-They may use httpx.AsyncClient against real FastAPI apps, real file I/O, etc.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-
-@pytest.fixture
-def anyio_backend():
- """Use asyncio as the async backend for integration tests."""
- return "asyncio"
diff --git a/tests/integration/rollout/__init__.py b/tests/integration/rollout/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/tests/integration/rollout/test_rollout_flow.py b/tests/integration/rollout/test_rollout_flow.py
deleted file mode 100644
index 721cf0dc..00000000
--- a/tests/integration/rollout/test_rollout_flow.py
+++ /dev/null
@@ -1,23 +0,0 @@
-"""Integration tests for the complete rollout lifecycle.
-
-Tests the full init -> chat -> completed flow using real FastAPI app
-with httpx.AsyncClient, minimal mocking (only the LLM client).
-"""
-
-from __future__ import annotations
-
-import pytest
-
-
-class TestRolloutLifecycle:
- """Test the complete rollout request lifecycle."""
-
- @pytest.mark.integration
- async def test_init_chat_complete_flow(self):
- """Full rollout: init request -> agent runs -> completion reported."""
- pytest.skip("TODO: implement rollout lifecycle integration test")
-
- @pytest.mark.integration
- async def test_concurrent_rollouts(self):
- """Multiple rollouts can execute concurrently without interference."""
- pytest.skip("TODO: implement concurrent rollout integration test")
diff --git a/tests/integration/server/__init__.py b/tests/integration/server/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/tests/integration/server/test_app_integration.py b/tests/integration/server/test_app_integration.py
deleted file mode 100644
index 2c40b6ab..00000000
--- a/tests/integration/server/test_app_integration.py
+++ /dev/null
@@ -1,23 +0,0 @@
-"""Integration tests for the FastAPI rollout server.
-
-Tests HTTP endpoints using httpx.AsyncClient against a real create_app() instance.
-Only the external LLM client is mocked.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-
-class TestServerEndpoints:
- """Test server HTTP endpoints with real request/response cycle."""
-
- @pytest.mark.integration
- async def test_health_endpoint(self):
- """GET /health returns 200 with server status."""
- pytest.skip("TODO: implement health endpoint integration test")
-
- @pytest.mark.integration
- async def test_rollout_endpoint_validation(self):
- """POST /v1/rollout/init validates request schema."""
- pytest.skip("TODO: implement rollout validation integration test")
From 0fc13b0c358c28dfcac0d65f882bdfa409bf6a57 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 18:10:14 -0800
Subject: [PATCH 21/60] Update test commands in CONTRIBUTING.md and enhance
concurrency control in tests
- Changed the test command in CONTRIBUTING.md to point to the correct test file.
- Improved the concurrency control test in test_app.py by implementing a ConcurrencyTrackingAgent to accurately track peak concurrent executions.
- Refactored the test_load_agent_loop_does_not_duplicate_cwd_in_sys_path to restore original sys.path after the test.
---
CONTRIBUTING.md | 2 +-
tests/unit/rollout/server/test_app.py | 38 +++++++++++++++++++--------
tests/unit/rollout/test_cli_utils.py | 5 +++-
3 files changed, 32 insertions(+), 13 deletions(-)
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index a8a0ce9c..754d452d 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -23,7 +23,7 @@ pre-commit install
pytest tests/
# Run a single test file
-pytest tests/test_rollout_base.py
+pytest tests/unit/rollout/core/test_base.py
# Run tests matching a pattern
pytest -k "test_name"
diff --git a/tests/unit/rollout/server/test_app.py b/tests/unit/rollout/server/test_app.py
index 909aa95f..c0c48c50 100644
--- a/tests/unit/rollout/server/test_app.py
+++ b/tests/unit/rollout/server/test_app.py
@@ -299,8 +299,30 @@ class TestConcurrencyControl:
async def test_semaphore_limits_concurrent_rollouts(self) -> None:
"""Only max_concurrent rollouts should run simultaneously."""
max_concurrent = 2
- agent = SlowAgentLoop(delay=0.3)
- app = create_app(agent, max_concurrent=max_concurrent)
+ peak_concurrent = 0
+ current_concurrent = 0
+ lock = asyncio.Lock()
+
+ class ConcurrencyTrackingAgent(RolloutAgentLoop):
+ name = "concurrency_agent"
+
+ def get_tools(
+ self, request: RolloutRequest
+ ) -> list[OpenAIFunctionToolSchema]:
+ return []
+
+ async def run(self, ctx: RolloutContext) -> RolloutResult:
+ nonlocal peak_concurrent, current_concurrent
+ async with lock:
+ current_concurrent += 1
+ if current_concurrent > peak_concurrent:
+ peak_concurrent = current_concurrent
+ await asyncio.sleep(0.3)
+ async with lock:
+ current_concurrent -= 1
+ return ctx.complete(list(ctx.request.messages))
+
+ app = create_app(ConcurrencyTrackingAgent(), max_concurrent=max_concurrent)
mock_client = mock_llm_client()
with patch("osmosis_ai.rollout.server.app.OsmosisLLMClient") as MockCls:
@@ -318,18 +340,12 @@ async def test_semaphore_limits_concurrent_rollouts(self) -> None:
)
assert resp.status_code == 202
- # Allow some time for background tasks
- await asyncio.sleep(0.1)
-
- # Check health to see active rollouts
- health = await client.get("/health")
- data = health.json()
- # There should be some active rollouts queued up
- assert data["active_rollouts"] >= 0
-
# Wait for all to complete
await asyncio.sleep(1.5)
+ # Peak concurrent executions must be capped by max_concurrent
+ assert peak_concurrent <= max_concurrent
+
async def test_max_concurrent_one_serializes_rollouts(self) -> None:
"""With max_concurrent=1 rollouts must execute one at a time."""
execution_order: list[str] = []
diff --git a/tests/unit/rollout/test_cli_utils.py b/tests/unit/rollout/test_cli_utils.py
index ad10a3e9..5f0fc51b 100644
--- a/tests/unit/rollout/test_cli_utils.py
+++ b/tests/unit/rollout/test_cli_utils.py
@@ -303,6 +303,9 @@ def test_load_agent_loop_does_not_duplicate_cwd_in_sys_path() -> None:
cwd = os.getcwd()
+ # Capture original sys.path so we can restore it
+ original_path = sys.path.copy()
+
# Ensure cwd IS in sys.path
if cwd not in sys.path:
sys.path.insert(0, cwd)
@@ -315,7 +318,7 @@ def test_load_agent_loop_does_not_duplicate_cwd_in_sys_path() -> None:
count_after = sys.path.count(cwd)
assert count_after == count_before
finally:
- pass
+ sys.path[:] = original_path
# =============================================================================
From 0c0dc62fbe3131aefaa916d119e9a775a6cfe617 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 19:55:11 -0800
Subject: [PATCH 22/60] Update pre-commit configuration to include additional
hooks and modify ruff hook
- Added new hooks for trailing whitespace, end-of-file fixing, YAML and TOML checks, merge conflict detection, and private key detection.
- Updated the ruff hook to use 'ruff-check' instead of 'ruff' with the '--fix' argument.
---
.pre-commit-config.yaml | 14 +++++++++++++-
1 file changed, 13 insertions(+), 1 deletion(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d7d2af79..e7bcb0f1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,7 +1,19 @@
repos:
+ - repo: https://github.com/pre-commit/pre-commit-hooks
+ rev: v6.0.0
+ hooks:
+ - id: trailing-whitespace
+ - id: end-of-file-fixer
+ - id: check-yaml
+ - id: check-toml
+ - id: check-merge-conflict
+ - id: check-added-large-files
+ args: ['--maxkb=500']
+ - id: detect-private-key
+
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.1
hooks:
- - id: ruff
+ - id: ruff-check
args: [--fix]
- id: ruff-format
From 35917fb8be3780c36a99765b3c06c4997ca45a21 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Mon, 16 Feb 2026 20:09:02 -0800
Subject: [PATCH 23/60] Update GitHub Actions workflow to enhance testing and
build processes
- Replaced ruff installation with the astral-sh/ruff-action for linting.
- Updated the test job to include Python versions 3.10, 3.11, 3.12, and 3.13.
- Modified the test execution steps to use uv for running tests and building the package.
- Added a new build job to check the package using uvx and twine.
---
.github/workflows/tests.yml | 59 +++++++++++++++++++++----------------
1 file changed, 34 insertions(+), 25 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 7edcd33f..ddfa44e5 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -9,42 +9,51 @@ on:
jobs:
lint:
runs-on: ubuntu-latest
-
steps:
- - name: Checkout repository
- uses: actions/checkout@v4
-
- - name: Set up Python
- uses: actions/setup-python@v5
- with:
- python-version: "3.12"
-
- - name: Install ruff
- run: pip install ruff==0.15.1
+ - uses: actions/checkout@v4
- - name: Ruff check
- run: ruff check .
+ - uses: astral-sh/ruff-action@v3
- - name: Ruff format check
- run: ruff format --check .
+ - uses: astral-sh/ruff-action@v3
+ with:
+ args: "format --check"
- pytest:
+ test:
runs-on: ubuntu-latest
strategy:
+ fail-fast: false
matrix:
- python-version: ["3.10", "3.12"]
-
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
steps:
- - name: Checkout repository
- uses: actions/checkout@v4
+ - uses: actions/checkout@v4
- - name: Set up Python ${{ matrix.python-version }}
- uses: actions/setup-python@v5
+ - uses: astral-sh/setup-uv@v7
with:
python-version: ${{ matrix.python-version }}
+ enable-cache: true
- name: Install dependencies
- run: pip install -e ".[dev]"
+ run: uv sync --frozen --extra dev
+
+ - name: Run tests with coverage
+ if: matrix.python-version == '3.12'
+ run: uv run pytest --cov=osmosis_ai --cov-report=term-missing
+
+ - name: Run tests
+ if: matrix.python-version != '3.12'
+ run: uv run pytest
+
+ build:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+
+ - uses: astral-sh/setup-uv@v7
+ with:
+ enable-cache: true
+
+ - name: Build package
+ run: uv build
- - name: Run pytest with coverage
- run: python -m pytest --cov=osmosis_ai --cov-report=term-missing
+ - name: Check package
+ run: uvx twine check --strict dist/*
From 99dcdd07946ea3dc54df46cee31c2bcc94c6df66 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 00:00:26 -0800
Subject: [PATCH 24/60] chore: add dual type checking with pyright and mypy
Introduce dedicated Pyright and Mypy checks in CI and align type annotations/configuration across rollout and evaluation modules. This catches type regressions earlier and improves typing reliability.
---
.github/workflows/tests.yml | 34 ++
CONTRIBUTING.md | 33 +-
osmosis_ai/_litellm_compat.py | 29 ++
osmosis_ai/cli_services/config.py | 9 +-
osmosis_ai/cli_services/engine.py | 9 +-
osmosis_ai/rollout/config/settings.py | 18 +-
osmosis_ai/rollout/console.py | 17 +-
osmosis_ai/rollout/core/base.py | 11 +
osmosis_ai/rollout/eval/common/cli.py | 2 +-
osmosis_ai/rollout/eval/common/errors.py | 8 +-
osmosis_ai/rollout/eval/common/llm_client.py | 45 +-
osmosis_ai/rollout/eval/common/runner.py | 2 +-
osmosis_ai/rollout/eval/test_mode/runner.py | 5 +-
osmosis_ai/rollout/server/app.py | 17 +-
osmosis_ai/rollout/testing.py | 11 +-
osmosis_ai/rubric_eval.py | 35 +-
osmosis_ai/utils.py | 38 +-
pyproject.toml | 26 +-
.../rollout/eval/common/test_llm_client.py | 235 +++-------
.../rollout/eval/evaluation/test_runner.py | 8 +-
tests/unit/rollout/mcp/test_agent_loop.py | 2 +-
tests/unit/test_rubric_eval.py | 441 ++++++++----------
uv.lock | 194 +++++++-
23 files changed, 744 insertions(+), 485 deletions(-)
create mode 100644 osmosis_ai/_litellm_compat.py
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index ddfa44e5..19797ebb 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -18,6 +18,40 @@ jobs:
with:
args: "format --check"
+ typecheck-pyright:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: astral-sh/setup-uv@v7
+ with:
+ python-version: "3.12"
+ enable-cache: true
+ - name: Install dependencies
+ run: uv sync --frozen --extra dev
+ - name: Run pyright
+ run: uv run pyright osmosis_ai/
+ - name: Build and install for verifytypes
+ run: |
+ uv build --wheel --out-dir dist/
+ uv pip install dist/*.whl
+ - name: Verify public API types
+ continue-on-error: true
+ run: uv run pyright --verifytypes osmosis_ai --ignoreexternal
+
+ typecheck-mypy:
+ runs-on: ubuntu-latest
+ continue-on-error: true
+ steps:
+ - uses: actions/checkout@v4
+ - uses: astral-sh/setup-uv@v7
+ with:
+ python-version: "3.12"
+ enable-cache: true
+ - name: Install dependencies
+ run: uv sync --frozen --extra dev
+ - name: Run mypy
+ run: uv run mypy osmosis_ai/
+
test:
runs-on: ubuntu-latest
strategy:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 754d452d..75cd08ea 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -6,11 +6,12 @@
git clone https://github.com/Osmosis-AI/osmosis-sdk-python
cd osmosis-sdk-python
-# Install with all development dependencies
+# Install with all development dependencies (includes server + mcp extras
+# for type checking and full test coverage)
pip install -e ".[dev]"
-# Or using uv
-uv sync --all-extras
+# Or using uv (recommended)
+uv sync --extra dev
# Install pre-commit hooks
pre-commit install
@@ -54,6 +55,30 @@ ruff format --check .
Ruff is pinned to one version across `pyproject.toml`, `.pre-commit-config.yaml`, and CI so local checks, pre-commit hooks, and GitHub Actions produce the same results.
+## Type Checking
+
+This project uses [Pyright](https://microsoft.github.io/pyright/) as the primary type checker and [mypy](https://mypy-lang.org/) as a secondary checker. Both are included in the `dev` extras — install with:
+
+```bash
+uv sync --extra dev
+```
+
+Then run via `uv run`:
+
+```bash
+# Pyright — must pass (blocking in CI)
+uv run pyright osmosis_ai/
+
+# mypy — advisory (non-blocking in CI)
+uv run mypy osmosis_ai/
+```
+
+Pyright runs in `standard` mode. All pyright errors must be resolved before merging. mypy runs in CI with `continue-on-error` and serves as an advisory check — fix mypy warnings when practical, but they won't block a PR.
+
+Configuration for both tools lives in `pyproject.toml` under `[tool.pyright]` and `[tool.mypy]`.
+
+> **Note:** CI also runs `pyright --verifytypes osmosis_ai --ignoreexternal` to check public API type completeness, but this command requires a non-editable install to locate the `py.typed` marker. It does not work with local editable installs (`uv sync`) and is therefore non-blocking in CI.
+
## Pre-commit Hooks
A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
@@ -72,4 +97,4 @@ This ensures `ruff check --fix` and `ruff format` run before each commit. Please
4. Run tests and linting
5. Submit a pull request
-CI will run linting (`ruff check` + `ruff format --check`) and tests with coverage on every PR.
+CI will run linting (`ruff check` + `ruff format --check`), type checking (pyright + mypy), and tests with coverage on every PR.
diff --git a/osmosis_ai/_litellm_compat.py b/osmosis_ai/_litellm_compat.py
new file mode 100644
index 00000000..1f52741a
--- /dev/null
+++ b/osmosis_ai/_litellm_compat.py
@@ -0,0 +1,29 @@
+# pyright: reportPrivateImportUsage=false
+"""Thin re-export layer for litellm types.
+
+pyright treats re-exports via a private-prefixed alias (_litellm) as private.
+We suppress the diagnostic at file level so the rest of the codebase can import
+cleanly.
+"""
+
+from __future__ import annotations
+
+import litellm as _litellm
+
+# ---------------------------------------------------------------------------
+# Exception types
+# ---------------------------------------------------------------------------
+NotFoundError = _litellm.NotFoundError
+APIError = _litellm.APIError
+RateLimitError = _litellm.RateLimitError
+AuthenticationError = _litellm.AuthenticationError
+Timeout = _litellm.Timeout
+APIConnectionError = _litellm.APIConnectionError
+BudgetExceededError = _litellm.BudgetExceededError
+ContextWindowExceededError = _litellm.ContextWindowExceededError
+
+# ---------------------------------------------------------------------------
+# Functions
+# ---------------------------------------------------------------------------
+completion = _litellm.completion
+acompletion = _litellm.acompletion
diff --git a/osmosis_ai/cli_services/config.py b/osmosis_ai/cli_services/config.py
index c94ea819..e0159780 100644
--- a/osmosis_ai/cli_services/config.py
+++ b/osmosis_ai/cli_services/config.py
@@ -4,11 +4,12 @@
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
-from typing import Any
+from typing import Any, cast
import yaml
from yaml.representer import SafeRepresenter
+from ..rubric_types import ModelInfo
from .errors import CLIError
from .shared import coerce_optional_float
@@ -23,7 +24,7 @@ class ParsedItem:
class RubricConfig:
rubric_id: str
rubric_text: str
- model_info: dict[str, Any]
+ model_info: ModelInfo
score_min: float | None
score_max: float | None
system_prompt: str | None
@@ -270,7 +271,7 @@ def _build_document_configs(
configs[rubric_key] = RubricConfig(
rubric_id=rubric_key,
rubric_text=rubric_text,
- model_info=copy.deepcopy(model_info),
+ model_info=copy.deepcopy(cast(ModelInfo, model_info)),
score_min=score_min,
score_max=score_max,
system_prompt=system_prompt if isinstance(system_prompt, str) else None,
@@ -422,7 +423,7 @@ class _LiteralSafeDumper(yaml.SafeDumper):
"""YAML dumper that preserves multiline strings with literal blocks."""
-def _represent_str(dumper: yaml.Dumper, data: str):
+def _represent_str(dumper: _LiteralSafeDumper, data: str): # type: ignore[type-arg]
if "\n" in data:
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
return SafeRepresenter.represent_str(dumper, data)
diff --git a/osmosis_ai/cli_services/engine.py b/osmosis_ai/cli_services/engine.py
index 9481b36b..eb834f1c 100644
--- a/osmosis_ai/cli_services/engine.py
+++ b/osmosis_ai/cli_services/engine.py
@@ -13,7 +13,12 @@
from tqdm import tqdm
from ..rubric_eval import DEFAULT_API_KEY_ENV, evaluate_rubric
-from ..rubric_types import MissingAPIKeyError, ModelNotFoundError, ProviderRequestError
+from ..rubric_types import (
+ MissingAPIKeyError,
+ ModelInfo,
+ ModelNotFoundError,
+ ProviderRequestError,
+)
from .config import RubricConfig
from .dataset import DatasetRecord
from .errors import CLIError
@@ -39,7 +44,7 @@ def _compose_extra_info_context(
api_key_env: str | None,
score_min: float | None,
score_max: float | None,
- model_info: dict[str, Any],
+ model_info: ModelInfo,
) -> tuple[dict[str, Any], dict[str, Any] | None]:
"""
Build the runtime context passed to rubric functions along with a sanitised
diff --git a/osmosis_ai/rollout/config/settings.py b/osmosis_ai/rollout/config/settings.py
index 40525372..aec5b2ee 100644
--- a/osmosis_ai/rollout/config/settings.py
+++ b/osmosis_ai/rollout/config/settings.py
@@ -19,21 +19,23 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from pydantic import BaseModel, Field
from osmosis_ai.rollout._compat import (
PYDANTIC_SETTINGS_AVAILABLE,
)
-# Conditionally use pydantic-settings BaseSettings or fallback to BaseModel
-if PYDANTIC_SETTINGS_AVAILABLE:
- from pydantic_settings import BaseSettings, SettingsConfigDict
-
- _BaseSettings = BaseSettings
+# Conditionally use pydantic-settings BaseSettings or fallback to BaseModel.
+# TYPE_CHECKING branch lets pyright/mypy resolve BaseSettings and its attributes;
+# the runtime branch handles the case where pydantic-settings is not installed.
+if TYPE_CHECKING or PYDANTIC_SETTINGS_AVAILABLE:
+ from pydantic_settings import BaseSettings as _BaseSettings
+ from pydantic_settings import SettingsConfigDict
else:
- # Fallback: use BaseModel (no env var loading)
- _BaseSettings = BaseModel # type: ignore[misc]
- SettingsConfigDict = None # type: ignore[misc, assignment]
+ _BaseSettings = BaseModel # type: ignore[assignment]
+ SettingsConfigDict = None # type: ignore[assignment]
class RolloutClientSettings(_BaseSettings):
diff --git a/osmosis_ai/rollout/console.py b/osmosis_ai/rollout/console.py
index d49331d7..c46c5c16 100644
--- a/osmosis_ai/rollout/console.py
+++ b/osmosis_ai/rollout/console.py
@@ -17,7 +17,14 @@
import sys
from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from rich import box as box
+ from rich.console import Console as RichConsole
+ from rich.markup import escape as rich_escape
+ from rich.panel import Panel
+ from rich.table import Table
# Try to import rich, gracefully degrade if not available
try:
@@ -25,19 +32,11 @@
from rich.console import Console as RichConsole
from rich.markup import escape as rich_escape
from rich.panel import Panel
- from rich.style import Style
from rich.table import Table
- from rich.text import Text
RICH_AVAILABLE = True
except ImportError:
RICH_AVAILABLE = False
- RichConsole = None # type: ignore
- Panel = None # type: ignore
- Table = None # type: ignore
- Text = None # type: ignore
- Style = None # type: ignore
- box = None # type: ignore
class _AnsiColors:
diff --git a/osmosis_ai/rollout/core/base.py b/osmosis_ai/rollout/core/base.py
index 789d8650..d39043d0 100644
--- a/osmosis_ai/rollout/core/base.py
+++ b/osmosis_ai/rollout/core/base.py
@@ -469,6 +469,17 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
"""
raise NotImplementedError
+ def get_default_tools(self) -> list[OpenAIFunctionToolSchema]:
+ """Return default tool list for discovery/validation (no real request needed)."""
+ return self.get_tools(
+ RolloutRequest(
+ rollout_id="discovery",
+ server_url="http://localhost",
+ messages=[],
+ completion_params={},
+ )
+ )
+
def __init_subclass__(cls, **kwargs: Any) -> None:
"""Validate that concrete subclasses define the name attribute."""
super().__init_subclass__(**kwargs)
diff --git a/osmosis_ai/rollout/eval/common/cli.py b/osmosis_ai/rollout/eval/common/cli.py
index 5d7cb550..0d2a985c 100644
--- a/osmosis_ai/rollout/eval/common/cli.py
+++ b/osmosis_ai/rollout/eval/common/cli.py
@@ -108,7 +108,7 @@ def load_mcp_agent(
agent_loop = MCPAgentLoop(mcp_server)
if not quiet:
- tool_names = [t.function.name for t in agent_loop.get_tools(None)] # type: ignore[arg-type]
+ tool_names = [t.function.name for t in agent_loop.get_default_tools()]
console.print(
f" Discovered {len(tool_names)} tool(s): {', '.join(tool_names)}"
)
diff --git a/osmosis_ai/rollout/eval/common/errors.py b/osmosis_ai/rollout/eval/common/errors.py
index 9aa45997..a4678dd8 100644
--- a/osmosis_ai/rollout/eval/common/errors.py
+++ b/osmosis_ai/rollout/eval/common/errors.py
@@ -42,7 +42,13 @@ class SystemicProviderError(ProviderError):
rather than retrying each row.
"""
- pass
+ duration_ms: float | None
+ tokens: int | None
+
+ def __init__(self, *args: object) -> None:
+ super().__init__(*args)
+ self.duration_ms = None
+ self.tokens = None
__all__ = [
diff --git a/osmosis_ai/rollout/eval/common/llm_client.py b/osmosis_ai/rollout/eval/common/llm_client.py
index 448335a8..fe7abfcc 100644
--- a/osmosis_ai/rollout/eval/common/llm_client.py
+++ b/osmosis_ai/rollout/eval/common/llm_client.py
@@ -17,6 +17,14 @@
category=UserWarning,
)
+from osmosis_ai._litellm_compat import APIConnectionError as _APIConnectionError
+from osmosis_ai._litellm_compat import AuthenticationError as _AuthenticationError
+from osmosis_ai._litellm_compat import BudgetExceededError as _BudgetExceededError
+from osmosis_ai._litellm_compat import (
+ ContextWindowExceededError as _ContextWindowExceededError,
+)
+from osmosis_ai._litellm_compat import RateLimitError as _RateLimitError
+from osmosis_ai._litellm_compat import Timeout as _LitellmTimeout
from osmosis_ai.rollout.client import CompletionsResult
from osmosis_ai.rollout.core.schemas import RolloutMetrics
from osmosis_ai.rollout.eval.common.errors import ProviderError, SystemicProviderError
@@ -141,12 +149,12 @@ def __init__(
) -> None:
self._litellm = _get_litellm()
- self._RateLimitError = self._litellm.RateLimitError
- self._AuthenticationError = self._litellm.AuthenticationError
- self._BudgetExceededError = self._litellm.BudgetExceededError
- self._Timeout = self._litellm.Timeout
- self._ContextWindowExceededError = self._litellm.ContextWindowExceededError
- self._APIConnectionError = self._litellm.APIConnectionError
+ self._RateLimitError = _RateLimitError
+ self._AuthenticationError = _AuthenticationError
+ self._BudgetExceededError = _BudgetExceededError
+ self._Timeout = _LitellmTimeout
+ self._ContextWindowExceededError = _ContextWindowExceededError
+ self._APIConnectionError = _APIConnectionError
# Preserve the user's original model name for display purposes.
self.display_name = model
@@ -251,14 +259,16 @@ async def chat_completions(
raise self._classify_unknown_error(e) from e
latency_ms = (time.monotonic() - start_time) * 1000
- usage = response.usage
+ # litellm's ModelResponse stubs are incomplete — usage/choices/message
+ # are valid at runtime but not fully reflected in the published types.
+ usage = response.usage # type: ignore[union-attr]
self._llm_latency_ms += latency_ms
self._num_llm_calls += 1
self._prompt_tokens += usage.prompt_tokens if usage else 0
self._response_tokens += usage.completion_tokens if usage else 0
- choice = response.choices[0]
- message = choice.message.model_dump(exclude_none=True)
+ choice = response.choices[0] # type: ignore[union-attr]
+ message = choice.message.model_dump(exclude_none=True) # type: ignore[union-attr]
return CompletionsResult(
message=message,
@@ -280,13 +290,16 @@ async def preflight_check(self) -> None:
fails, or the model does not exist.
"""
try:
- await self._litellm.acompletion(
- model=self.model,
- messages=[{"role": "user", "content": "hi"}],
- max_tokens=1,
- **({"api_key": self._api_key} if self._api_key else {}),
- **({"api_base": self._api_base} if self._api_base else {}),
- )
+ preflight_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": [{"role": "user", "content": "hi"}],
+ "max_tokens": 1,
+ }
+ if self._api_key:
+ preflight_kwargs["api_key"] = self._api_key
+ if self._api_base:
+ preflight_kwargs["api_base"] = self._api_base
+ await self._litellm.acompletion(**preflight_kwargs)
except self._RateLimitError:
# Rate-limited means the endpoint is reachable and authenticated.
return
diff --git a/osmosis_ai/rollout/eval/common/runner.py b/osmosis_ai/rollout/eval/common/runner.py
index 346f91a3..c831946f 100644
--- a/osmosis_ai/rollout/eval/common/runner.py
+++ b/osmosis_ai/rollout/eval/common/runner.py
@@ -241,7 +241,7 @@ async def run_batch(
duration_ms = getattr(e, "duration_ms", None)
if duration_ms is None:
duration_ms = (time.monotonic() - row_start) * 1000
- tokens = int(getattr(e, "tokens", 0))
+ tokens = int(e.tokens or 0)
result = LocalRunResult(
row_index=row_index,
success=False,
diff --git a/osmosis_ai/rollout/eval/test_mode/runner.py b/osmosis_ai/rollout/eval/test_mode/runner.py
index 80286520..d3d4418a 100644
--- a/osmosis_ai/rollout/eval/test_mode/runner.py
+++ b/osmosis_ai/rollout/eval/test_mode/runner.py
@@ -55,13 +55,16 @@ async def run_single(
row_index: int,
max_turns: int = 10,
completion_params: dict[str, Any] | None = None,
+ rollout_id: str | None = None,
+ request_metadata: dict[str, Any] | None = None,
) -> LocalTestRunResult:
return await super().run_single(
row=row,
row_index=row_index,
max_turns=max_turns,
completion_params=completion_params,
- rollout_id=f"test-{row_index}",
+ rollout_id=rollout_id or f"test-{row_index}",
+ request_metadata=request_metadata,
)
diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py
index 1cbf9de1..a0bd38fe 100644
--- a/osmosis_ai/rollout/server/app.py
+++ b/osmosis_ai/rollout/server/app.py
@@ -180,6 +180,11 @@ async def lifespan(app: FastAPI):
register_with_platform,
)
+ # Bind narrowed values to local variables for the closure
+ _reg_host: str = server_host
+ _reg_port: int = server_port
+ _reg_credentials: WorkspaceCredentials = credentials
+
async def do_registration():
import httpx
@@ -188,7 +193,7 @@ async def do_registration():
state.settings.registration_readiness_poll_interval_seconds
)
timeout = state.settings.registration_readiness_timeout_seconds
- health_url = f"http://127.0.0.1:{server_port}/health"
+ health_url = f"http://127.0.0.1:{_reg_port}/health"
start_time = time.monotonic()
server_ready = False
@@ -220,16 +225,16 @@ async def do_registration():
# Run sync registration in thread pool to avoid blocking event loop
result = await asyncio.to_thread(
register_with_platform,
- host=server_host,
- port=server_port,
+ host=_reg_host,
+ port=_reg_port,
agent_loop_name=agent_loop.name,
- credentials=credentials,
+ credentials=_reg_credentials,
api_key=api_key,
)
print_registration_result(
result=result,
- host=server_host,
- port=server_port,
+ host=_reg_host,
+ port=_reg_port,
agent_loop_name=agent_loop.name,
api_key=api_key,
)
diff --git a/osmosis_ai/rollout/testing.py b/osmosis_ai/rollout/testing.py
index 1bdc3829..e24de716 100644
--- a/osmosis_ai/rollout/testing.py
+++ b/osmosis_ai/rollout/testing.py
@@ -34,6 +34,7 @@
from fastapi.testclient import TestClient
from osmosis_ai.rollout.core.schemas import (
+ CompletionsChoice,
CompletionsRequest,
CompletionsResponse,
CompletionUsage,
@@ -231,11 +232,11 @@ async def completions(request: CompletionsRequest) -> CompletionsResponse:
created=int(time.time()),
model=request.model,
choices=[
- {
- "index": 0,
- "message": assistant_message,
- "finish_reason": "stop",
- }
+ CompletionsChoice(
+ index=0,
+ message=assistant_message,
+ finish_reason="stop",
+ )
],
usage=CompletionUsage(
prompt_tokens=len(prompt_token_ids),
diff --git a/osmosis_ai/rubric_eval.py b/osmosis_ai/rubric_eval.py
index 5c493fac..352bf044 100644
--- a/osmosis_ai/rubric_eval.py
+++ b/osmosis_ai/rubric_eval.py
@@ -15,6 +15,27 @@
import warnings
from typing import Any
+from ._litellm_compat import (
+ APIConnectionError as _LitellmAPIConnectionError,
+)
+from ._litellm_compat import (
+ APIError as _LitellmAPIError,
+)
+from ._litellm_compat import (
+ AuthenticationError as _LitellmAuthenticationError,
+)
+from ._litellm_compat import (
+ NotFoundError as _LitellmNotFoundError,
+)
+from ._litellm_compat import (
+ RateLimitError as _LitellmRateLimitError,
+)
+from ._litellm_compat import (
+ Timeout as _LitellmTimeout,
+)
+from ._litellm_compat import (
+ completion as _litellm_completion,
+)
from .rubric_types import (
MissingAPIKeyError,
ModelInfo,
@@ -423,8 +444,8 @@ def _call_litellm(
)
try:
- response = litellm.completion(**completion_kwargs)
- except litellm.NotFoundError as err:
+ response = _litellm_completion(**completion_kwargs)
+ except _LitellmNotFoundError as err:
raise ModelNotFoundError(
provider,
model,
@@ -432,11 +453,11 @@ def _call_litellm(
f"and your {provider} account has access to it.",
) from err
except (
- litellm.APIError,
- litellm.RateLimitError,
- litellm.AuthenticationError,
- litellm.Timeout,
- litellm.APIConnectionError,
+ _LitellmAPIError,
+ _LitellmRateLimitError,
+ _LitellmAuthenticationError,
+ _LitellmTimeout,
+ _LitellmAPIConnectionError,
) as err:
raise ProviderRequestError(
provider, model, _extract_error_message(err)
diff --git a/osmosis_ai/utils.py b/osmosis_ai/utils.py
index d36b51ff..5945da80 100644
--- a/osmosis_ai/utils.py
+++ b/osmosis_ai/utils.py
@@ -95,19 +95,20 @@ def _check_classic_return(result):
if asyncio.iscoroutinefunction(func):
@functools.wraps(func)
- async def wrapper(*args, **kwargs):
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
if _drop_data_source:
kwargs.pop("data_source", None)
return _check_classic_return(await func(*args, **kwargs))
- else:
- @functools.wraps(func)
- def wrapper(*args, **kwargs):
- if _drop_data_source:
- kwargs.pop("data_source", None)
- return _check_classic_return(func(*args, **kwargs))
+ return wrapper
- return wrapper
+ @functools.wraps(func)
+ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+ if _drop_data_source:
+ kwargs.pop("data_source", None)
+ return _check_classic_return(func(*args, **kwargs))
+
+ return sync_wrapper
def _is_str_annotation(annotation: Any) -> bool:
@@ -262,20 +263,21 @@ def _check_classic_return(result):
if asyncio.iscoroutinefunction(func):
@functools.wraps(func)
- async def wrapper(*args, **kwargs):
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
if _drop_data_source:
kwargs.pop("data_source", None)
if is_classic:
_validate_classic_args(*args, **kwargs)
return _check_classic_return(await func(*args, **kwargs))
- else:
- @functools.wraps(func)
- def wrapper(*args, **kwargs):
- if _drop_data_source:
- kwargs.pop("data_source", None)
- if is_classic:
- _validate_classic_args(*args, **kwargs)
- return _check_classic_return(func(*args, **kwargs))
+ return wrapper
+
+ @functools.wraps(func)
+ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+ if _drop_data_source:
+ kwargs.pop("data_source", None)
+ if is_classic:
+ _validate_classic_args(*args, **kwargs)
+ return _check_classic_return(func(*args, **kwargs))
- return wrapper
+ return sync_wrapper
diff --git a/pyproject.toml b/pyproject.toml
index 74d3fc53..59f1c577 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,15 +43,18 @@ mcp = [
"fastmcp>=2.0.0",
]
-# Development dependencies (testing, formatting, etc.)
+# Development dependencies (testing, formatting, type checking, etc.)
+# Includes server + mcp extras so pyright/mypy can resolve all imports.
dev = [
+ "osmosis-ai[server,mcp]",
"pytest>=8.0.0,<10.0.0",
"pytest-asyncio>=0.23.0,<2.0.0",
"pytest-cov>=6.0.0",
"anyio>=4.0.0",
"ruff==0.15.1", # Keep in sync with pre-commit and CI
"pre-commit>=4.0.0,<5.0.0",
- "fastmcp>=2.0.0",
+ "pyright[nodejs]==1.1.408",
+ "mypy==1.19.1",
]
# Full installation with all optional features (union of server + mcp)
@@ -123,3 +126,22 @@ ignore = [
[tool.ruff.lint.isort]
known-first-party = ["osmosis_ai"]
+
+[tool.pyright]
+pythonVersion = "3.10"
+include = ["osmosis_ai"]
+exclude = ["tests", "build", "dist", "**/__pycache__"]
+typeCheckingMode = "standard"
+reportMissingTypeStubs = false
+reportCallInDefaultInitializer = false
+
+[tool.mypy]
+python_version = "3.10"
+files = ["osmosis_ai"]
+exclude = ["tests/", "build/", "dist/"]
+check_untyped_defs = true
+warn_return_any = true
+warn_unused_configs = true
+ignore_missing_imports = true
+warn_unused_ignores = false
+disallow_incomplete_defs = true
diff --git a/tests/unit/rollout/eval/common/test_llm_client.py b/tests/unit/rollout/eval/common/test_llm_client.py
index 3f3d4ad0..ea6ea876 100644
--- a/tests/unit/rollout/eval/common/test_llm_client.py
+++ b/tests/unit/rollout/eval/common/test_llm_client.py
@@ -8,10 +8,57 @@
import pytest
+from osmosis_ai._litellm_compat import (
+ APIConnectionError as _APIConnectionError,
+)
+from osmosis_ai._litellm_compat import (
+ APIError as _APIError,
+)
+from osmosis_ai._litellm_compat import (
+ AuthenticationError as _AuthenticationError,
+)
+from osmosis_ai._litellm_compat import (
+ BudgetExceededError as _BudgetExceededError,
+)
+from osmosis_ai._litellm_compat import (
+ ContextWindowExceededError as _ContextWindowExceededError,
+)
+from osmosis_ai._litellm_compat import (
+ RateLimitError as _RateLimitError,
+)
+from osmosis_ai._litellm_compat import (
+ Timeout as _LitellmTimeout,
+)
from osmosis_ai.rollout.client import CompletionsResult
from osmosis_ai.rollout.core.schemas import OpenAIFunctionToolSchema
+def _make_litellm_error(cls: type, message: str, **attrs: Any) -> Exception:
+ """Instantiate a real litellm exception with required positional args.
+
+ litellm exception constructors require varying positional parameters
+ (message, model, llm_provider, etc.). This helper fills them in so
+ tests only need to specify the human-readable *message*.
+ """
+ if cls is _BudgetExceededError:
+ exc = cls(current_cost=100.0, max_budget=50.0, message=message)
+ elif cls is _APIError:
+ exc = cls(
+ status_code=attrs.pop("status_code", 500),
+ message=message,
+ llm_provider="openai",
+ model="gpt-4o",
+ )
+ elif cls is _LitellmTimeout or cls is _ContextWindowExceededError:
+ exc = cls(message=message, model="gpt-4o", llm_provider="openai")
+ else:
+ # RateLimitError, AuthenticationError, APIConnectionError
+ exc = cls(message=message, llm_provider="openai", model="gpt-4o")
+ for k, v in attrs.items():
+ setattr(exc, k, v)
+ return exc
+
+
class TestExternalLLMClient:
"""Tests for ExternalLLMClient."""
@@ -267,16 +314,6 @@ async def capture_kwargs(**kwargs):
async def test_wraps_missing_provider_error_with_compact_hint(self) -> None:
"""Provider inference failures should return concise actionable errors."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
mock_litellm.acompletion = AsyncMock(
side_effect=Exception(
"litellm.BadRequestError: LLM Provider NOT provided. "
@@ -314,16 +351,6 @@ class TestPreflightCheck:
async def test_preflight_success(self) -> None:
"""Preflight succeeds when acompletion returns normally."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
mock_response = MagicMock()
mock_response.choices = [
MagicMock(
@@ -354,18 +381,7 @@ async def test_preflight_success(self) -> None:
async def test_preflight_auth_failure(self) -> None:
"""Preflight raises SystemicProviderError on auth failure."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- auth_error = mock_litellm.AuthenticationError("Invalid API key")
+ auth_error = _make_litellm_error(_AuthenticationError, "Invalid API key")
mock_litellm.acompletion = AsyncMock(side_effect=auth_error)
with patch(
@@ -383,18 +399,7 @@ async def test_preflight_auth_failure(self) -> None:
async def test_preflight_connection_failure(self) -> None:
"""Preflight raises SystemicProviderError on connection failure."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- conn_error = mock_litellm.APIConnectionError("Connection refused")
+ conn_error = _make_litellm_error(_APIConnectionError, "Connection refused")
mock_litellm.acompletion = AsyncMock(side_effect=conn_error)
with patch(
@@ -412,18 +417,7 @@ async def test_preflight_connection_failure(self) -> None:
async def test_preflight_rate_limit_passes(self) -> None:
"""Preflight treats RateLimitError as success (endpoint is reachable)."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- rate_error = mock_litellm.RateLimitError("Rate limited")
+ rate_error = _make_litellm_error(_RateLimitError, "Rate limited")
mock_litellm.acompletion = AsyncMock(side_effect=rate_error)
with patch(
@@ -440,19 +434,10 @@ async def test_preflight_rate_limit_passes(self) -> None:
async def test_preflight_404_api_error(self) -> None:
"""Preflight raises SystemicProviderError on 404 API error."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- api_error = mock_litellm.APIError("Model not found")
- api_error.status_code = 404
+ # Use a generic Exception with status_code=404 so it falls through
+ # to _classify_unknown_error which checks status_code.
+ api_error = Exception("Model not found")
+ api_error.status_code = 404 # type: ignore[attr-defined]
mock_litellm.acompletion = AsyncMock(side_effect=api_error)
with patch(
@@ -472,23 +457,14 @@ async def test_preflight_connection_error_wrapped_as_internal_server_error(
) -> None:
"""Preflight catches connection errors wrapped as InternalServerError (APIError subclass)."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
- # litellm wraps connection errors as InternalServerError (APIError subclass)
- InternalServerError = type("InternalServerError", (mock_litellm.APIError,), {})
- api_error = InternalServerError(
+ # Simulate litellm wrapping connection errors as InternalServerError.
+ # Use a plain Exception with status_code and connection-error message
+ # so it falls through to _classify_unknown_error.
+ api_error = Exception(
"InternalServerError: OpenAIException - Connection error."
)
- api_error.status_code = 500
+ api_error.status_code = 500 # type: ignore[attr-defined]
mock_litellm.acompletion = AsyncMock(side_effect=api_error)
with patch(
@@ -513,19 +489,8 @@ class TestSystemicErrorClassification:
async def test_auth_error_is_systemic(self) -> None:
"""AuthenticationError should raise SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
mock_litellm.acompletion = AsyncMock(
- side_effect=mock_litellm.AuthenticationError("Invalid key")
+ side_effect=_make_litellm_error(_AuthenticationError, "Invalid key")
)
with patch(
@@ -543,19 +508,8 @@ async def test_auth_error_is_systemic(self) -> None:
async def test_connection_error_is_systemic(self) -> None:
"""APIConnectionError should raise SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
mock_litellm.acompletion = AsyncMock(
- side_effect=mock_litellm.APIConnectionError("Connection refused")
+ side_effect=_make_litellm_error(_APIConnectionError, "Connection refused")
)
with patch(
@@ -573,19 +527,8 @@ async def test_connection_error_is_systemic(self) -> None:
async def test_rate_limit_is_not_systemic(self) -> None:
"""RateLimitError should raise ProviderError, not SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
mock_litellm.acompletion = AsyncMock(
- side_effect=mock_litellm.RateLimitError("Too many requests")
+ side_effect=_make_litellm_error(_RateLimitError, "Too many requests")
)
with patch(
@@ -609,19 +552,10 @@ async def test_rate_limit_is_not_systemic(self) -> None:
async def test_api_error_401_is_systemic(self) -> None:
"""APIError with 401 status should raise SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- api_error = mock_litellm.APIError("Unauthorized")
- api_error.status_code = 401
+ # Use a generic Exception with status_code=401 so it falls through
+ # to _classify_unknown_error which checks status_code.
+ api_error = Exception("Unauthorized")
+ api_error.status_code = 401 # type: ignore[attr-defined]
mock_litellm.acompletion = AsyncMock(side_effect=api_error)
with patch(
@@ -639,23 +573,13 @@ async def test_api_error_401_is_systemic(self) -> None:
async def test_connection_error_wrapped_as_api_error_is_systemic(self) -> None:
"""APIError with 'Connection error' message should raise SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- # Simulate litellm wrapping connection error as InternalServerError
- InternalServerError = type("InternalServerError", (mock_litellm.APIError,), {})
- api_error = InternalServerError(
+ # Simulate litellm wrapping connection error as InternalServerError.
+ # Use a plain Exception with status_code and connection-error message
+ # so it falls through to _classify_unknown_error.
+ api_error = Exception(
"InternalServerError: OpenAIException - Connection error."
)
- api_error.status_code = 500
+ api_error.status_code = 500 # type: ignore[attr-defined]
mock_litellm.acompletion = AsyncMock(side_effect=api_error)
with patch(
@@ -676,19 +600,10 @@ async def test_connection_error_wrapped_as_api_error_is_systemic(self) -> None:
async def test_api_error_500_is_not_systemic(self) -> None:
"""APIError with 500 status should raise ProviderError, not SystemicProviderError."""
mock_litellm = MagicMock()
- for error_name in (
- "RateLimitError",
- "AuthenticationError",
- "APIError",
- "BudgetExceededError",
- "Timeout",
- "ContextWindowExceededError",
- "APIConnectionError",
- ):
- setattr(mock_litellm, error_name, type(error_name, (Exception,), {}))
-
- api_error = mock_litellm.APIError("Internal error")
- api_error.status_code = 500
+ # Use a generic Exception with status_code=500 so it falls through
+ # to _classify_unknown_error which checks status_code.
+ api_error = Exception("Internal error")
+ api_error.status_code = 500 # type: ignore[attr-defined]
mock_litellm.acompletion = AsyncMock(side_effect=api_error)
with patch(
diff --git a/tests/unit/rollout/eval/evaluation/test_runner.py b/tests/unit/rollout/eval/evaluation/test_runner.py
index b4e077e9..a3008ecd 100644
--- a/tests/unit/rollout/eval/evaluation/test_runner.py
+++ b/tests/unit/rollout/eval/evaluation/test_runner.py
@@ -901,8 +901,8 @@ def test_uses_attributes_when_present(self) -> None:
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
e = SystemicProviderError("test")
- e.duration_ms = 1234.5 # type: ignore[attr-defined]
- e.tokens = 42 # type: ignore[attr-defined]
+ e.duration_ms = 1234.5
+ e.tokens = 42
dur, tok = _extract_systemic_error_metrics(e, fallback_started_at=0.0)
assert dur == 1234.5
assert tok == 42
@@ -924,8 +924,8 @@ def test_fallback_tokens_when_none(self) -> None:
from osmosis_ai.rollout.eval.common.errors import SystemicProviderError
e = SystemicProviderError("test")
- e.duration_ms = 100.0 # type: ignore[attr-defined]
- e.tokens = None # type: ignore[attr-defined]
+ e.duration_ms = 100.0
+ e.tokens = None
dur, tok = _extract_systemic_error_metrics(e, fallback_started_at=0.0)
assert dur == 100.0
assert tok == 0
diff --git a/tests/unit/rollout/mcp/test_agent_loop.py b/tests/unit/rollout/mcp/test_agent_loop.py
index 25bc848d..0b6be21b 100644
--- a/tests/unit/rollout/mcp/test_agent_loop.py
+++ b/tests/unit/rollout/mcp/test_agent_loop.py
@@ -225,7 +225,7 @@ def test_returns_converted_schemas(self):
)
agent = MCPAgentLoop(server)
- tools = agent.get_tools(None) # type: ignore[arg-type]
+ tools = agent.get_default_tools()
assert len(tools) == 2
names = {t.function.name for t in tools}
diff --git a/tests/unit/test_rubric_eval.py b/tests/unit/test_rubric_eval.py
index be03e66c..b046ad09 100644
--- a/tests/unit/test_rubric_eval.py
+++ b/tests/unit/test_rubric_eval.py
@@ -1,4 +1,4 @@
-"""Tests for osmosis_ai.rubric_eval — prompt building, JSON parsing, and LiteLLM evaluation."""
+"""Tests for osmosis_ai.rubric_eval -- prompt building, JSON parsing, and LiteLLM evaluation."""
import json
import sys
@@ -224,145 +224,169 @@ def _create_mock_litellm_response(score: float, explanation: str) -> MagicMock:
return mock_response
+# The _litellm_completion reference is captured at import time from the compat
+# layer, so we must patch it on the rubric_eval module directly rather than
+# replacing sys.modules["litellm"]. The litellm module-level import inside
+# _call_litellm (used for suppress_debug_info and supports_response_schema)
+# is still patched via sys.modules.
+_COMPLETION_PATCH = "osmosis_ai.rubric_eval._litellm_completion"
+
+
class TestEvaluateRubric:
"""Tests for the main evaluate_rubric function."""
- def test_evaluate_rubric_success(self):
+ @pytest.fixture()
+ def mock_rubric_litellm(self):
+ """Mock both litellm module and _litellm_completion for rubric eval tests."""
mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
+ mock_completion = MagicMock()
+ with (
+ patch.dict(sys.modules, {"litellm": mock_litellm}),
+ patch(_COMPLETION_PATCH, mock_completion),
+ ):
+ yield mock_litellm, mock_completion
+
+ def test_evaluate_rubric_success(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(
0.85, "Good response"
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score accuracy",
- solution_str="The answer is 42",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
+ result = evaluate_rubric(
+ rubric="Score accuracy",
+ solution_str="The answer is 42",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
assert result == 0.85
- mock_litellm.completion.assert_called_once()
- call_kwargs = mock_litellm.completion.call_args.kwargs
+ mock_completion.assert_called_once()
+ call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["model"] == "openai/gpt-4o"
assert call_kwargs["api_key"] == "test-key"
assert call_kwargs["temperature"] == 0
- def test_evaluate_rubric_gpt5_no_temperature(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.9, "Excellent"
+ def test_evaluate_rubric_gpt5_no_temperature(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(0.9, "Excellent")
+
+ evaluate_rubric(
+ rubric="Score quality",
+ solution_str="Test response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5-mini",
+ "api_key": "test-key",
+ },
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- evaluate_rubric(
- rubric="Score quality",
- solution_str="Test response",
- model_info={
- "provider": "openai",
- "model": "gpt-5-mini",
- "api_key": "test-key",
- },
- )
-
- call_kwargs = mock_litellm.completion.call_args.kwargs
+ call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["model"] == "openai/responses/gpt-5-mini"
assert "temperature" not in call_kwargs
- def test_evaluate_rubric_anthropic(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.9, "Excellent"
+ def test_evaluate_rubric_anthropic(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(0.9, "Excellent")
+
+ result = evaluate_rubric(
+ rubric="Score quality",
+ solution_str="Well written response",
+ model_info={
+ "provider": "anthropic",
+ "model": "claude-sonnet-4-5-20250929",
+ "api_key": "test-key",
+ },
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score quality",
- solution_str="Well written response",
- model_info={
- "provider": "anthropic",
- "model": "claude-sonnet-4-5-20250929",
- "api_key": "test-key",
- },
- )
-
assert result == 0.9
- call_kwargs = mock_litellm.completion.call_args.kwargs
+ call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["model"] == "anthropic/claude-sonnet-4-5-20250929"
- def test_evaluate_rubric_xai(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 0.8, "Good"
+ def test_evaluate_rubric_xai(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(0.8, "Good")
+
+ result = evaluate_rubric(
+ rubric="Score response",
+ solution_str="Some response",
+ model_info={
+ "provider": "xai",
+ "model": "grok-4-fast",
+ "api_key": "test-key",
+ },
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score response",
- solution_str="Some response",
- model_info={
- "provider": "xai",
- "model": "grok-4-fast",
- "api_key": "test-key",
- },
- )
-
assert result == 0.8
- call_kwargs = mock_litellm.completion.call_args.kwargs
+ call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["model"] == "xai/grok-4-fast"
- def test_evaluate_rubric_return_details(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
+ def test_evaluate_rubric_return_details(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(
0.65, "Detailed explanation"
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score accuracy",
- solution_str="Some answer",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- return_details=True,
- )
+ result = evaluate_rubric(
+ rubric="Score accuracy",
+ solution_str="Some answer",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ return_details=True,
+ )
assert result["score"] == 0.65
assert result["explanation"] == "Detailed explanation"
assert "raw" in result
- def test_evaluate_rubric_score_clamping_max(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 1.5, "Over max"
+ def test_evaluate_rubric_score_clamping_max(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(1.5, "Over max")
+
+ result = evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
assert result == 1.0 # Clamped to max
- def test_evaluate_rubric_score_clamping_min(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- -0.5, "Under min"
+ def test_evaluate_rubric_score_clamping_min(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(-0.5, "Under min")
+
+ result = evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
)
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
+ assert result == 0.0 # Clamped to min
+
+ def test_evaluate_rubric_api_error(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.side_effect = litellm.APIError(
+ message="API rate limit exceeded",
+ llm_provider="openai",
+ model="gpt-4o",
+ status_code=429,
+ )
+
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
rubric="Score",
solution_str="Response",
model_info={
@@ -372,184 +396,133 @@ def test_evaluate_rubric_score_clamping_min(self):
},
)
- assert result == 0.0 # Clamped to min
-
- def test_evaluate_rubric_api_error(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.APIError(
- message="API rate limit exceeded",
- llm_provider="openai",
- model="gpt-4o",
- status_code=429,
- )
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
-
assert "rate limit" in str(exc_info.value).lower()
- def test_evaluate_rubric_invalid_json_response(self):
- mock_litellm = MagicMock()
+ def test_evaluate_rubric_invalid_json_response(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
mock_response = MagicMock()
mock_response.model_dump.return_value = {
"choices": [{"message": {"content": "Not valid JSON"}}]
}
- mock_litellm.completion.return_value = mock_response
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
+ mock_completion.return_value = mock_response
+
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
assert "not valid JSON" in str(exc_info.value)
- def test_evaluate_rubric_empty_response(self):
- mock_litellm = MagicMock()
+ def test_evaluate_rubric_empty_response(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
mock_response = MagicMock()
mock_response.model_dump.return_value = {
"choices": [{"message": {"content": ""}}]
}
- mock_litellm.completion.return_value = mock_response
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
+ mock_completion.return_value = mock_response
+
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ )
assert "did not include any content" in str(exc_info.value)
- def test_evaluate_rubric_model_not_found_404(self):
+ def test_evaluate_rubric_model_not_found_404(self, mock_rubric_litellm):
"""NotFoundError from litellm should raise ModelNotFoundError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.NotFoundError(
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.side_effect = litellm.NotFoundError(
message="The model `no-such-model` does not exist",
llm_provider="openai",
model="no-such-model",
)
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ModelNotFoundError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "no-such-model",
- "api_key": "test-key",
- },
- )
+
+ with pytest.raises(ModelNotFoundError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "no-such-model",
+ "api_key": "test-key",
+ },
+ )
assert exc_info.value.provider == "openai"
assert exc_info.value.model == "no-such-model"
- def test_evaluate_rubric_model_not_found_is_provider_request_error(self):
+ def test_evaluate_rubric_model_not_found_is_provider_request_error(
+ self, mock_rubric_litellm
+ ):
"""ModelNotFoundError should be catchable as ProviderRequestError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.NotFoundError(
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.side_effect = litellm.NotFoundError(
message="Not found",
llm_provider="anthropic",
model="bad-model",
)
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError):
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "anthropic",
- "model": "bad-model",
- "api_key": "test-key",
- },
- )
-
- def test_evaluate_rubric_non_404_error_is_not_model_not_found(self):
+
+ with pytest.raises(ProviderRequestError):
+ evaluate_rubric(
+ rubric="Score",
+ solution_str="Response",
+ model_info={
+ "provider": "anthropic",
+ "model": "bad-model",
+ "api_key": "test-key",
+ },
+ )
+
+ def test_evaluate_rubric_non_404_error_is_not_model_not_found(
+ self, mock_rubric_litellm
+ ):
"""Non-NotFoundError should raise ProviderRequestError, not ModelNotFoundError."""
- mock_litellm = MagicMock()
- mock_litellm.completion.side_effect = litellm.RateLimitError(
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.side_effect = litellm.RateLimitError(
message="Rate limit exceeded",
llm_provider="openai",
model="gpt-4o",
)
- mock_litellm.APIError = litellm.APIError
- mock_litellm.NotFoundError = litellm.NotFoundError
- mock_litellm.RateLimitError = litellm.RateLimitError
- mock_litellm.AuthenticationError = litellm.AuthenticationError
- mock_litellm.Timeout = litellm.Timeout
- mock_litellm.APIConnectionError = litellm.APIConnectionError
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- with pytest.raises(ProviderRequestError) as exc_info:
- evaluate_rubric(
- rubric="Score",
- solution_str="Response",
- model_info={
- "provider": "openai",
- "model": "gpt-4o",
- "api_key": "test-key",
- },
- )
- assert not isinstance(exc_info.value, ModelNotFoundError)
-
- def test_evaluate_rubric_custom_score_range(self):
- mock_litellm = MagicMock()
- mock_litellm.completion.return_value = _create_mock_litellm_response(
- 7.5, "Good"
- )
-
- with patch.dict(sys.modules, {"litellm": mock_litellm}):
- result = evaluate_rubric(
- rubric="Score from 0-10",
+ with pytest.raises(ProviderRequestError) as exc_info:
+ evaluate_rubric(
+ rubric="Score",
solution_str="Response",
model_info={
"provider": "openai",
"model": "gpt-4o",
"api_key": "test-key",
},
- score_min=0.0,
- score_max=10.0,
)
+ assert not isinstance(exc_info.value, ModelNotFoundError)
+
+ def test_evaluate_rubric_custom_score_range(self, mock_rubric_litellm):
+ _, mock_completion = mock_rubric_litellm
+ mock_completion.return_value = _create_mock_litellm_response(7.5, "Good")
+
+ result = evaluate_rubric(
+ rubric="Score from 0-10",
+ solution_str="Response",
+ model_info={
+ "provider": "openai",
+ "model": "gpt-4o",
+ "api_key": "test-key",
+ },
+ score_min=0.0,
+ score_max=10.0,
+ )
+
assert result == 7.5
diff --git a/uv.lock b/uv.lock
index 74424eba..0c632111 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1436,6 +1436,91 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" },
]
+[[package]]
+name = "librt"
+version = "0.8.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/8a/3f/4ca7dd7819bf8ff303aca39c3c60e5320e46e766ab7f7dd627d3b9c11bdf/librt-0.8.0.tar.gz", hash = "sha256:cb74cdcbc0103fc988e04e5c58b0b31e8e5dd2babb9182b6f9490488eb36324b", size = 177306, upload-time = "2026-02-12T14:53:54.743Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d5/e9/018cfd60629e0404e6917943789800aa2231defbea540a17b90cc4547b97/librt-0.8.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:db63cf3586a24241e89ca1ce0b56baaec9d371a328bd186c529b27c914c9a1ef", size = 65690, upload-time = "2026-02-12T14:51:57.761Z" },
+ { url = "https://files.pythonhosted.org/packages/b5/80/8d39980860e4d1c9497ee50e5cd7c4766d8cfd90d105578eae418e8ffcbc/librt-0.8.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ba9d9e60651615bc614be5e21a82cdb7b1769a029369cf4b4d861e4f19686fb6", size = 68373, upload-time = "2026-02-12T14:51:59.013Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/76/6e6f7a443af63977e421bd542551fec4072d9eaba02e671b05b238fe73bc/librt-0.8.0-cp310-cp310-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:cb4b3ad543084ed79f186741470b251b9d269cd8b03556f15a8d1a99a64b7de5", size = 197091, upload-time = "2026-02-12T14:52:00.642Z" },
+ { url = "https://files.pythonhosted.org/packages/14/40/fa064181c231334c9f4cb69eb338132d39510c8928e84beba34b861d0a71/librt-0.8.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3d2720335020219197380ccfa5c895f079ac364b4c429e96952cd6509934d8eb", size = 207350, upload-time = "2026-02-12T14:52:02.32Z" },
+ { url = "https://files.pythonhosted.org/packages/50/49/e7f8438dd226305e3e5955d495114ad01448e6a6ffc0303289b4153b5fc5/librt-0.8.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9726305d3e53419d27fc8cdfcd3f9571f0ceae22fa6b5ea1b3662c2e538f833e", size = 219962, upload-time = "2026-02-12T14:52:03.884Z" },
+ { url = "https://files.pythonhosted.org/packages/1f/2c/74086fc5d52e77107a3cc80a9a3209be6ad1c9b6bc99969d8d9bbf9fdfe4/librt-0.8.0-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:cc3d107f603b5ee7a79b6aa6f166551b99b32fb4a5303c4dfcb4222fc6a0335e", size = 212939, upload-time = "2026-02-12T14:52:05.537Z" },
+ { url = "https://files.pythonhosted.org/packages/c8/ae/d6917c0ebec9bc2e0293903d6a5ccc7cdb64c228e529e96520b277318f25/librt-0.8.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:41064a0c07b4cc7a81355ccc305cb097d6027002209ffca51306e65ee8293630", size = 221393, upload-time = "2026-02-12T14:52:07.164Z" },
+ { url = "https://files.pythonhosted.org/packages/04/97/15df8270f524ce09ad5c19cbbe0e8f95067582507149a6c90594e7795370/librt-0.8.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:c6e4c10761ddbc0d67d2f6e2753daf99908db85d8b901729bf2bf5eaa60e0567", size = 216721, upload-time = "2026-02-12T14:52:08.857Z" },
+ { url = "https://files.pythonhosted.org/packages/c4/52/17cbcf9b7a1bae5016d9d3561bc7169b32c3bd216c47d934d3f270602c0c/librt-0.8.0-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:ba581acad5ac8f33e2ff1746e8a57e001b47c6721873121bf8bbcf7ba8bd3aa4", size = 214790, upload-time = "2026-02-12T14:52:10.033Z" },
+ { url = "https://files.pythonhosted.org/packages/2a/2d/010a236e8dc4d717dd545c46fd036dcced2c7ede71ef85cf55325809ff92/librt-0.8.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:bdab762e2c0b48bab76f1a08acb3f4c77afd2123bedac59446aeaaeed3d086cf", size = 237384, upload-time = "2026-02-12T14:52:11.244Z" },
+ { url = "https://files.pythonhosted.org/packages/38/14/f1c0eff3df8760dee761029efb72991c554d9f3282f1048e8c3d0eb60997/librt-0.8.0-cp310-cp310-win32.whl", hash = "sha256:6a3146c63220d814c4a2c7d6a1eacc8d5c14aed0ff85115c1dfea868080cd18f", size = 54289, upload-time = "2026-02-12T14:52:12.798Z" },
+ { url = "https://files.pythonhosted.org/packages/2f/0b/2684d473e64890882729f91866ed97ccc0a751a0afc3b4bf1a7b57094dbb/librt-0.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:bbebd2bba5c6ae02907df49150e55870fdd7440d727b6192c46b6f754723dde9", size = 61347, upload-time = "2026-02-12T14:52:13.793Z" },
+ { url = "https://files.pythonhosted.org/packages/51/e9/42af181c89b65abfd557c1b017cba5b82098eef7bf26d1649d82ce93ccc7/librt-0.8.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0ce33a9778e294507f3a0e3468eccb6a698b5166df7db85661543eca1cfc5369", size = 65314, upload-time = "2026-02-12T14:52:14.778Z" },
+ { url = "https://files.pythonhosted.org/packages/9d/4a/15a847fca119dc0334a4b8012b1e15fdc5fc19d505b71e227eaf1bcdba09/librt-0.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8070aa3368559de81061ef752770d03ca1f5fc9467d4d512d405bd0483bfffe6", size = 68015, upload-time = "2026-02-12T14:52:15.797Z" },
+ { url = "https://files.pythonhosted.org/packages/e1/87/ffc8dbd6ab68dd91b736c88529411a6729649d2b74b887f91f3aaff8d992/librt-0.8.0-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:20f73d4fecba969efc15cdefd030e382502d56bb6f1fc66b580cce582836c9fa", size = 194508, upload-time = "2026-02-12T14:52:16.835Z" },
+ { url = "https://files.pythonhosted.org/packages/89/92/a7355cea28d6c48ff6ff5083ac4a2a866fb9b07b786aa70d1f1116680cd5/librt-0.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a512c88900bdb1d448882f5623a0b1ad27ba81a9bd75dacfe17080b72272ca1f", size = 205630, upload-time = "2026-02-12T14:52:18.58Z" },
+ { url = "https://files.pythonhosted.org/packages/ac/5e/54509038d7ac527828db95b8ba1c8f5d2649bc32fd8f39b1718ec9957dce/librt-0.8.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:015e2dde6e096d27c10238bf9f6492ba6c65822dfb69d2bf74c41a8e88b7ddef", size = 218289, upload-time = "2026-02-12T14:52:20.134Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/17/0ee0d13685cefee6d6f2d47bb643ddad3c62387e2882139794e6a5f1288a/librt-0.8.0-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1c25a131013eadd3c600686a0c0333eb2896483cbc7f65baa6a7ee761017aef9", size = 211508, upload-time = "2026-02-12T14:52:21.413Z" },
+ { url = "https://files.pythonhosted.org/packages/4b/a8/1714ef6e9325582e3727de3be27e4c1b2f428ea411d09f1396374180f130/librt-0.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:21b14464bee0b604d80a638cf1ee3148d84ca4cc163dcdcecb46060c1b3605e4", size = 219129, upload-time = "2026-02-12T14:52:22.61Z" },
+ { url = "https://files.pythonhosted.org/packages/89/d3/2d9fe353edff91cdc0ece179348054a6fa61f3de992c44b9477cb973509b/librt-0.8.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:05a3dd3f116747f7e1a2b475ccdc6fb637fd4987126d109e03013a79d40bf9e6", size = 213126, upload-time = "2026-02-12T14:52:23.819Z" },
+ { url = "https://files.pythonhosted.org/packages/ad/8e/9f5c60444880f6ad50e3ff7475e5529e787797e7f3ad5432241633733b92/librt-0.8.0-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:fa37f99bff354ff191c6bcdffbc9d7cdd4fc37faccfc9be0ef3a4fd5613977da", size = 212279, upload-time = "2026-02-12T14:52:25.034Z" },
+ { url = "https://files.pythonhosted.org/packages/fe/eb/d4a2cfa647da3022ae977f50d7eda1d91f70d7d1883cf958a4b6ef689eab/librt-0.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1566dbb9d1eb0987264c9b9460d212e809ba908d2f4a3999383a84d765f2f3f1", size = 234654, upload-time = "2026-02-12T14:52:26.204Z" },
+ { url = "https://files.pythonhosted.org/packages/6a/31/26b978861c7983b036a3aea08bdbb2ec32bbaab1ad1d57c5e022be59afc1/librt-0.8.0-cp311-cp311-win32.whl", hash = "sha256:70defb797c4d5402166787a6b3c66dfb3fa7f93d118c0509ffafa35a392f4258", size = 54603, upload-time = "2026-02-12T14:52:27.342Z" },
+ { url = "https://files.pythonhosted.org/packages/d0/78/f194ed7c48dacf875677e749c5d0d1d69a9daa7c994314a39466237fb1be/librt-0.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:db953b675079884ffda33d1dca7189fb961b6d372153750beb81880384300817", size = 61730, upload-time = "2026-02-12T14:52:28.31Z" },
+ { url = "https://files.pythonhosted.org/packages/97/ee/ad71095478d02137b6f49469dc808c595cfe89b50985f6b39c5345f0faab/librt-0.8.0-cp311-cp311-win_arm64.whl", hash = "sha256:75d1a8cab20b2043f03f7aab730551e9e440adc034d776f15f6f8d582b0a5ad4", size = 52274, upload-time = "2026-02-12T14:52:29.345Z" },
+ { url = "https://files.pythonhosted.org/packages/fb/53/f3bc0c4921adb0d4a5afa0656f2c0fbe20e18e3e0295e12985b9a5dc3f55/librt-0.8.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:17269dd2745dbe8e42475acb28e419ad92dfa38214224b1b01020b8cac70b645", size = 66511, upload-time = "2026-02-12T14:52:30.34Z" },
+ { url = "https://files.pythonhosted.org/packages/89/4b/4c96357432007c25a1b5e363045373a6c39481e49f6ba05234bb59a839c1/librt-0.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f4617cef654fca552f00ce5ffdf4f4b68770f18950e4246ce94629b789b92467", size = 68628, upload-time = "2026-02-12T14:52:31.491Z" },
+ { url = "https://files.pythonhosted.org/packages/47/16/52d75374d1012e8fc709216b5eaa25f471370e2a2331b8be00f18670a6c7/librt-0.8.0-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5cb11061a736a9db45e3c1293cfcb1e3caf205912dfa085734ba750f2197ff9a", size = 198941, upload-time = "2026-02-12T14:52:32.489Z" },
+ { url = "https://files.pythonhosted.org/packages/fc/11/d5dd89e5a2228567b1228d8602d896736247424484db086eea6b8010bcba/librt-0.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4bb00bd71b448f16749909b08a0ff16f58b079e2261c2e1000f2bbb2a4f0a45", size = 210009, upload-time = "2026-02-12T14:52:33.634Z" },
+ { url = "https://files.pythonhosted.org/packages/49/d8/fc1a92a77c3020ee08ce2dc48aed4b42ab7c30fb43ce488d388673b0f164/librt-0.8.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:95a719a049f0eefaf1952673223cf00d442952273cbd20cf2ed7ec423a0ef58d", size = 224461, upload-time = "2026-02-12T14:52:34.868Z" },
+ { url = "https://files.pythonhosted.org/packages/7f/98/eb923e8b028cece924c246104aa800cf72e02d023a8ad4ca87135b05a2fe/librt-0.8.0-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bd32add59b58fba3439d48d6f36ac695830388e3da3e92e4fc26d2d02670d19c", size = 217538, upload-time = "2026-02-12T14:52:36.078Z" },
+ { url = "https://files.pythonhosted.org/packages/fd/67/24e80ab170674a1d8ee9f9a83081dca4635519dbd0473b8321deecddb5be/librt-0.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4f764b2424cb04524ff7a486b9c391e93f93dc1bd8305b2136d25e582e99aa2f", size = 225110, upload-time = "2026-02-12T14:52:37.301Z" },
+ { url = "https://files.pythonhosted.org/packages/d8/c7/6fbdcbd1a6e5243c7989c21d68ab967c153b391351174b4729e359d9977f/librt-0.8.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:f04ca50e847abc486fa8f4107250566441e693779a5374ba211e96e238f298b9", size = 217758, upload-time = "2026-02-12T14:52:38.89Z" },
+ { url = "https://files.pythonhosted.org/packages/4b/bd/4d6b36669db086e3d747434430073e14def032dd58ad97959bf7e2d06c67/librt-0.8.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:9ab3a3475a55b89b87ffd7e6665838e8458e0b596c22e0177e0f961434ec474a", size = 218384, upload-time = "2026-02-12T14:52:40.637Z" },
+ { url = "https://files.pythonhosted.org/packages/50/2d/afe966beb0a8f179b132f3e95c8dd90738a23e9ebdba10f89a3f192f9366/librt-0.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3e36a8da17134ffc29373775d88c04832f9ecfab1880470661813e6c7991ef79", size = 241187, upload-time = "2026-02-12T14:52:43.55Z" },
+ { url = "https://files.pythonhosted.org/packages/02/d0/6172ea4af2b538462785ab1a68e52d5c99cfb9866a7caf00fdf388299734/librt-0.8.0-cp312-cp312-win32.whl", hash = "sha256:4eb5e06ebcc668677ed6389164f52f13f71737fc8be471101fa8b4ce77baeb0c", size = 54914, upload-time = "2026-02-12T14:52:44.676Z" },
+ { url = "https://files.pythonhosted.org/packages/d4/cb/ceb6ed6175612a4337ad49fb01ef594712b934b4bc88ce8a63554832eb44/librt-0.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:0a33335eb59921e77c9acc05d0e654e4e32e45b014a4d61517897c11591094f8", size = 62020, upload-time = "2026-02-12T14:52:45.676Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/7e/61701acbc67da74ce06ddc7ba9483e81c70f44236b2d00f6a4bfee1aacbf/librt-0.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:24a01c13a2a9bdad20997a4443ebe6e329df063d1978bbe2ebbf637878a46d1e", size = 52443, upload-time = "2026-02-12T14:52:47.218Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/32/3edb0bcb4113a9c8bdcd1750663a54565d255027657a5df9d90f13ee07fa/librt-0.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7f820210e21e3a8bf8fde2ae3c3d10106d4de9ead28cbfdf6d0f0f41f5b12fa1", size = 66522, upload-time = "2026-02-12T14:52:48.219Z" },
+ { url = "https://files.pythonhosted.org/packages/30/ab/e8c3d05e281f5d405ebdcc5bc8ab36df23e1a4b40ac9da8c3eb9928b72b9/librt-0.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4831c44b8919e75ca0dfb52052897c1ef59fdae19d3589893fbd068f1e41afbf", size = 68658, upload-time = "2026-02-12T14:52:50.351Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/d3/74a206c47b7748bbc8c43942de3ed67de4c231156e148b4f9250869593df/librt-0.8.0-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:88c6e75540f1f10f5e0fc5e87b4b6c290f0e90d1db8c6734f670840494764af8", size = 199287, upload-time = "2026-02-12T14:52:51.938Z" },
+ { url = "https://files.pythonhosted.org/packages/fa/29/ef98a9131cf12cb95771d24e4c411fda96c89dc78b09c2de4704877ebee4/librt-0.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9646178cd794704d722306c2c920c221abbf080fede3ba539d5afdec16c46dad", size = 210293, upload-time = "2026-02-12T14:52:53.128Z" },
+ { url = "https://files.pythonhosted.org/packages/5b/3e/89b4968cb08c53d4c2d8b02517081dfe4b9e07a959ec143d333d76899f6c/librt-0.8.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e1af31a710e17891d9adf0dbd9a5fcd94901a3922a96499abdbf7ce658f4e01", size = 224801, upload-time = "2026-02-12T14:52:54.367Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/28/f38526d501f9513f8b48d78e6be4a241e15dd4b000056dc8b3f06ee9ce5d/librt-0.8.0-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:507e94f4bec00b2f590fbe55f48cd518a208e2474a3b90a60aa8f29136ddbada", size = 218090, upload-time = "2026-02-12T14:52:55.758Z" },
+ { url = "https://files.pythonhosted.org/packages/02/ec/64e29887c5009c24dc9c397116c680caffc50286f62bd99c39e3875a2854/librt-0.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f1178e0de0c271231a660fbef9be6acdfa1d596803464706862bef6644cc1cae", size = 225483, upload-time = "2026-02-12T14:52:57.375Z" },
+ { url = "https://files.pythonhosted.org/packages/ee/16/7850bdbc9f1a32d3feff2708d90c56fc0490b13f1012e438532781aa598c/librt-0.8.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:71fc517efc14f75c2f74b1f0a5d5eb4a8e06aa135c34d18eaf3522f4a53cd62d", size = 218226, upload-time = "2026-02-12T14:52:58.534Z" },
+ { url = "https://files.pythonhosted.org/packages/1c/4a/166bffc992d65ddefa7c47052010a87c059b44a458ebaf8f5eba384b0533/librt-0.8.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:0583aef7e9a720dd40f26a2ad5a1bf2ccbb90059dac2b32ac516df232c701db3", size = 218755, upload-time = "2026-02-12T14:52:59.701Z" },
+ { url = "https://files.pythonhosted.org/packages/da/5d/9aeee038bcc72a9cfaaee934463fe9280a73c5440d36bd3175069d2cb97b/librt-0.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5d0f76fc73480d42285c609c0ea74d79856c160fa828ff9aceab574ea4ecfd7b", size = 241617, upload-time = "2026-02-12T14:53:00.966Z" },
+ { url = "https://files.pythonhosted.org/packages/64/ff/2bec6b0296b9d0402aa6ec8540aa19ebcb875d669c37800cb43d10d9c3a3/librt-0.8.0-cp313-cp313-win32.whl", hash = "sha256:e79dbc8f57de360f0ed987dc7de7be814b4803ef0e8fc6d3ff86e16798c99935", size = 54966, upload-time = "2026-02-12T14:53:02.042Z" },
+ { url = "https://files.pythonhosted.org/packages/08/8d/bf44633b0182996b2c7ea69a03a5c529683fa1f6b8e45c03fe874ff40d56/librt-0.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:25b3e667cbfc9000c4740b282df599ebd91dbdcc1aa6785050e4c1d6be5329ab", size = 62000, upload-time = "2026-02-12T14:53:03.822Z" },
+ { url = "https://files.pythonhosted.org/packages/5c/fd/c6472b8e0eac0925001f75e366cf5500bcb975357a65ef1f6b5749389d3a/librt-0.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:e9a3a38eb4134ad33122a6d575e6324831f930a771d951a15ce232e0237412c2", size = 52496, upload-time = "2026-02-12T14:53:04.889Z" },
+ { url = "https://files.pythonhosted.org/packages/e0/13/79ebfe30cd273d7c0ce37a5f14dc489c5fb8b722a008983db2cfd57270bb/librt-0.8.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:421765e8c6b18e64d21c8ead315708a56fc24f44075059702e421d164575fdda", size = 66078, upload-time = "2026-02-12T14:53:06.085Z" },
+ { url = "https://files.pythonhosted.org/packages/4b/8f/d11eca40b62a8d5e759239a80636386ef88adecb10d1a050b38cc0da9f9e/librt-0.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:48f84830a8f8ad7918afd743fd7c4eb558728bceab7b0e38fd5a5cf78206a556", size = 68309, upload-time = "2026-02-12T14:53:07.121Z" },
+ { url = "https://files.pythonhosted.org/packages/9c/b4/f12ee70a3596db40ff3c88ec9eaa4e323f3b92f77505b4d900746706ec6a/librt-0.8.0-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:9f09d4884f882baa39a7e36bbf3eae124c4ca2a223efb91e567381d1c55c6b06", size = 196804, upload-time = "2026-02-12T14:53:08.164Z" },
+ { url = "https://files.pythonhosted.org/packages/8b/7e/70dbbdc0271fd626abe1671ad117bcd61a9a88cdc6a10ccfbfc703db1873/librt-0.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:693697133c3b32aa9b27f040e3691be210e9ac4d905061859a9ed519b1d5a376", size = 206915, upload-time = "2026-02-12T14:53:09.333Z" },
+ { url = "https://files.pythonhosted.org/packages/79/13/6b9e05a635d4327608d06b3c1702166e3b3e78315846373446cf90d7b0bf/librt-0.8.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c5512aae4648152abaf4d48b59890503fcbe86e85abc12fb9b096fe948bdd816", size = 221200, upload-time = "2026-02-12T14:53:10.68Z" },
+ { url = "https://files.pythonhosted.org/packages/35/6c/e19a3ac53e9414de43a73d7507d2d766cd22d8ca763d29a4e072d628db42/librt-0.8.0-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:995d24caa6bbb34bcdd4a41df98ac6d1af637cfa8975cb0790e47d6623e70e3e", size = 214640, upload-time = "2026-02-12T14:53:12.342Z" },
+ { url = "https://files.pythonhosted.org/packages/30/f0/23a78464788619e8c70f090cfd099cce4973eed142c4dccb99fc322283fd/librt-0.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b9aef96d7593584e31ef6ac1eb9775355b0099fee7651fae3a15bc8657b67b52", size = 221980, upload-time = "2026-02-12T14:53:13.603Z" },
+ { url = "https://files.pythonhosted.org/packages/03/32/38e21420c5d7aa8a8bd2c7a7d5252ab174a5a8aaec8b5551968979b747bf/librt-0.8.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:4f6e975377fbc4c9567cb33ea9ab826031b6c7ec0515bfae66a4fb110d40d6da", size = 215146, upload-time = "2026-02-12T14:53:14.8Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/00/bd9ecf38b1824c25240b3ad982fb62c80f0a969e6679091ba2b3afb2b510/librt-0.8.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:daae5e955764be8fd70a93e9e5133c75297f8bce1e802e1d3683b98f77e1c5ab", size = 215203, upload-time = "2026-02-12T14:53:16.087Z" },
+ { url = "https://files.pythonhosted.org/packages/b9/60/7559bcc5279d37810b98d4a52616febd7b8eef04391714fd6bdf629598b1/librt-0.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:7bd68cebf3131bb920d5984f75fe302d758db33264e44b45ad139385662d7bc3", size = 237937, upload-time = "2026-02-12T14:53:17.236Z" },
+ { url = "https://files.pythonhosted.org/packages/41/cc/be3e7da88f1abbe2642672af1dc00a0bccece11ca60241b1883f3018d8d5/librt-0.8.0-cp314-cp314-win32.whl", hash = "sha256:1e6811cac1dcb27ca4c74e0ca4a5917a8e06db0d8408d30daee3a41724bfde7a", size = 50685, upload-time = "2026-02-12T14:53:18.888Z" },
+ { url = "https://files.pythonhosted.org/packages/38/27/e381d0df182a8f61ef1f6025d8b138b3318cc9d18ad4d5f47c3bf7492523/librt-0.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:178707cda89d910c3b28bf5aa5f69d3d4734e0f6ae102f753ad79edef83a83c7", size = 57872, upload-time = "2026-02-12T14:53:19.942Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/0c/ca9dfdf00554a44dea7d555001248269a4bab569e1590a91391feb863fa4/librt-0.8.0-cp314-cp314-win_arm64.whl", hash = "sha256:3e8b77b5f54d0937b26512774916041756c9eb3e66f1031971e626eea49d0bf4", size = 48056, upload-time = "2026-02-12T14:53:21.473Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/ed/6cc9c4ad24f90c8e782193c7b4a857408fd49540800613d1356c63567d7b/librt-0.8.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:789911e8fa40a2e82f41120c936b1965f3213c67f5a483fc5a41f5839a05dcbb", size = 68307, upload-time = "2026-02-12T14:53:22.498Z" },
+ { url = "https://files.pythonhosted.org/packages/84/d8/0e94292c6b3e00b6eeea39dd44d5703d1ec29b6dafce7eea19dc8f1aedbd/librt-0.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2b37437e7e4ef5e15a297b36ba9e577f73e29564131d86dd75875705e97402b5", size = 70999, upload-time = "2026-02-12T14:53:23.603Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/f4/6be1afcbdeedbdbbf54a7c9d73ad43e1bf36897cebf3978308cd64922e02/librt-0.8.0-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:671a6152edf3b924d98a5ed5e6982ec9cb30894085482acadce0975f031d4c5c", size = 220782, upload-time = "2026-02-12T14:53:25.133Z" },
+ { url = "https://files.pythonhosted.org/packages/f0/8d/f306e8caa93cfaf5c6c9e0d940908d75dc6af4fd856baa5535c922ee02b1/librt-0.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8992ca186a1678107b0af3d0c9303d8c7305981b9914989b9788319ed4d89546", size = 235420, upload-time = "2026-02-12T14:53:27.047Z" },
+ { url = "https://files.pythonhosted.org/packages/d6/f2/65d86bd462e9c351326564ca805e8457442149f348496e25ccd94583ffa2/librt-0.8.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:001e5330093d887b8b9165823eca6c5c4db183fe4edea4fdc0680bbac5f46944", size = 246452, upload-time = "2026-02-12T14:53:28.341Z" },
+ { url = "https://files.pythonhosted.org/packages/03/94/39c88b503b4cb3fcbdeb3caa29672b6b44ebee8dcc8a54d49839ac280f3f/librt-0.8.0-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d920789eca7ef71df7f31fd547ec0d3002e04d77f30ba6881e08a630e7b2c30e", size = 238891, upload-time = "2026-02-12T14:53:29.625Z" },
+ { url = "https://files.pythonhosted.org/packages/e3/c6/6c0d68190893d01b71b9569b07a1c811e280c0065a791249921c83dc0290/librt-0.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:82fb4602d1b3e303a58bfe6165992b5a78d823ec646445356c332cd5f5bbaa61", size = 250249, upload-time = "2026-02-12T14:53:30.93Z" },
+ { url = "https://files.pythonhosted.org/packages/52/7a/f715ed9e039035d0ea637579c3c0155ab3709a7046bc408c0fb05d337121/librt-0.8.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:4d3e38797eb482485b486898f89415a6ab163bc291476bd95712e42cf4383c05", size = 240642, upload-time = "2026-02-12T14:53:32.174Z" },
+ { url = "https://files.pythonhosted.org/packages/c2/3c/609000a333debf5992efe087edc6467c1fdbdddca5b610355569bbea9589/librt-0.8.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:a905091a13e0884701226860836d0386b88c72ce5c2fdfba6618e14c72be9f25", size = 239621, upload-time = "2026-02-12T14:53:33.39Z" },
+ { url = "https://files.pythonhosted.org/packages/b9/df/87b0673d5c395a8f34f38569c116c93142d4dc7e04af2510620772d6bd4f/librt-0.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:375eda7acfce1f15f5ed56cfc960669eefa1ec8732e3e9087c3c4c3f2066759c", size = 262986, upload-time = "2026-02-12T14:53:34.617Z" },
+ { url = "https://files.pythonhosted.org/packages/09/7f/6bbbe9dcda649684773aaea78b87fff4d7e59550fbc2877faa83612087a3/librt-0.8.0-cp314-cp314t-win32.whl", hash = "sha256:2ccdd20d9a72c562ffb73098ac411de351b53a6fbb3390903b2d33078ef90447", size = 51328, upload-time = "2026-02-12T14:53:36.15Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/f3/e1981ab6fa9b41be0396648b5850267888a752d025313a9e929c4856208e/librt-0.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:25e82d920d4d62ad741592fcf8d0f3bda0e3fc388a184cb7d2f566c681c5f7b9", size = 58719, upload-time = "2026-02-12T14:53:37.183Z" },
+ { url = "https://files.pythonhosted.org/packages/94/d1/433b3c06e78f23486fe4fdd19bc134657eb30997d2054b0dbf52bbf3382e/librt-0.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:92249938ab744a5890580d3cb2b22042f0dce71cdaa7c1369823df62bedf7cbc", size = 48753, upload-time = "2026-02-12T14:53:38.539Z" },
+]
+
[[package]]
name = "litellm"
version = "1.80.15"
@@ -1812,6 +1897,61 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/da/7d22601b625e241d4f23ef1ebff8acfc60da633c9e7e7922e24d10f592b3/multidict-6.7.0-py3-none-any.whl", hash = "sha256:394fc5c42a333c9ffc3e421a4c85e08580d990e08b99f6bf35b4132114c5dcb3", size = 12317, upload-time = "2025-10-06T14:52:29.272Z" },
]
+[[package]]
+name = "mypy"
+version = "1.19.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "librt", marker = "platform_python_implementation != 'PyPy'" },
+ { name = "mypy-extensions" },
+ { name = "pathspec" },
+ { name = "tomli", marker = "python_full_version < '3.11'" },
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f5/db/4efed9504bc01309ab9c2da7e352cc223569f05478012b5d9ece38fd44d2/mypy-1.19.1.tar.gz", hash = "sha256:19d88bb05303fe63f71dd2c6270daca27cb9401c4ca8255fe50d1d920e0eb9ba", size = 3582404, upload-time = "2025-12-15T05:03:48.42Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/2f/63/e499890d8e39b1ff2df4c0c6ce5d371b6844ee22b8250687a99fd2f657a8/mypy-1.19.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5f05aa3d375b385734388e844bc01733bd33c644ab48e9684faa54e5389775ec", size = 13101333, upload-time = "2025-12-15T05:03:03.28Z" },
+ { url = "https://files.pythonhosted.org/packages/72/4b/095626fc136fba96effc4fd4a82b41d688ab92124f8c4f7564bffe5cf1b0/mypy-1.19.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:022ea7279374af1a5d78dfcab853fe6a536eebfda4b59deab53cd21f6cd9f00b", size = 12164102, upload-time = "2025-12-15T05:02:33.611Z" },
+ { url = "https://files.pythonhosted.org/packages/0c/5b/952928dd081bf88a83a5ccd49aaecfcd18fd0d2710c7ff07b8fb6f7032b9/mypy-1.19.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee4c11e460685c3e0c64a4c5de82ae143622410950d6be863303a1c4ba0e36d6", size = 12765799, upload-time = "2025-12-15T05:03:28.44Z" },
+ { url = "https://files.pythonhosted.org/packages/2a/0d/93c2e4a287f74ef11a66fb6d49c7a9f05e47b0a4399040e6719b57f500d2/mypy-1.19.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de759aafbae8763283b2ee5869c7255391fbc4de3ff171f8f030b5ec48381b74", size = 13522149, upload-time = "2025-12-15T05:02:36.011Z" },
+ { url = "https://files.pythonhosted.org/packages/7b/0e/33a294b56aaad2b338d203e3a1d8b453637ac36cb278b45005e0901cf148/mypy-1.19.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:ab43590f9cd5108f41aacf9fca31841142c786827a74ab7cc8a2eacb634e09a1", size = 13810105, upload-time = "2025-12-15T05:02:40.327Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/fd/3e82603a0cb66b67c5e7abababce6bf1a929ddf67bf445e652684af5c5a0/mypy-1.19.1-cp310-cp310-win_amd64.whl", hash = "sha256:2899753e2f61e571b3971747e302d5f420c3fd09650e1951e99f823bc3089dac", size = 10057200, upload-time = "2025-12-15T05:02:51.012Z" },
+ { url = "https://files.pythonhosted.org/packages/ef/47/6b3ebabd5474d9cdc170d1342fbf9dddc1b0ec13ec90bf9004ee6f391c31/mypy-1.19.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d8dfc6ab58ca7dda47d9237349157500468e404b17213d44fc1cb77bce532288", size = 13028539, upload-time = "2025-12-15T05:03:44.129Z" },
+ { url = "https://files.pythonhosted.org/packages/5c/a6/ac7c7a88a3c9c54334f53a941b765e6ec6c4ebd65d3fe8cdcfbe0d0fd7db/mypy-1.19.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e3f276d8493c3c97930e354b2595a44a21348b320d859fb4a2b9f66da9ed27ab", size = 12083163, upload-time = "2025-12-15T05:03:37.679Z" },
+ { url = "https://files.pythonhosted.org/packages/67/af/3afa9cf880aa4a2c803798ac24f1d11ef72a0c8079689fac5cfd815e2830/mypy-1.19.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2abb24cf3f17864770d18d673c85235ba52456b36a06b6afc1e07c1fdcd3d0e6", size = 12687629, upload-time = "2025-12-15T05:02:31.526Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/46/20f8a7114a56484ab268b0ab372461cb3a8f7deed31ea96b83a4e4cfcfca/mypy-1.19.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a009ffa5a621762d0c926a078c2d639104becab69e79538a494bcccb62cc0331", size = 13436933, upload-time = "2025-12-15T05:03:15.606Z" },
+ { url = "https://files.pythonhosted.org/packages/5b/f8/33b291ea85050a21f15da910002460f1f445f8007adb29230f0adea279cb/mypy-1.19.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f7cee03c9a2e2ee26ec07479f38ea9c884e301d42c6d43a19d20fb014e3ba925", size = 13661754, upload-time = "2025-12-15T05:02:26.731Z" },
+ { url = "https://files.pythonhosted.org/packages/fd/a3/47cbd4e85bec4335a9cd80cf67dbc02be21b5d4c9c23ad6b95d6c5196bac/mypy-1.19.1-cp311-cp311-win_amd64.whl", hash = "sha256:4b84a7a18f41e167f7995200a1d07a4a6810e89d29859df936f1c3923d263042", size = 10055772, upload-time = "2025-12-15T05:03:26.179Z" },
+ { url = "https://files.pythonhosted.org/packages/06/8a/19bfae96f6615aa8a0604915512e0289b1fad33d5909bf7244f02935d33a/mypy-1.19.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a8174a03289288c1f6c46d55cef02379b478bfbc8e358e02047487cad44c6ca1", size = 13206053, upload-time = "2025-12-15T05:03:46.622Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/34/3e63879ab041602154ba2a9f99817bb0c85c4df19a23a1443c8986e4d565/mypy-1.19.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffcebe56eb09ff0c0885e750036a095e23793ba6c2e894e7e63f6d89ad51f22e", size = 12219134, upload-time = "2025-12-15T05:03:24.367Z" },
+ { url = "https://files.pythonhosted.org/packages/89/cc/2db6f0e95366b630364e09845672dbee0cbf0bbe753a204b29a944967cd9/mypy-1.19.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b64d987153888790bcdb03a6473d321820597ab8dd9243b27a92153c4fa50fd2", size = 12731616, upload-time = "2025-12-15T05:02:44.725Z" },
+ { url = "https://files.pythonhosted.org/packages/00/be/dd56c1fd4807bc1eba1cf18b2a850d0de7bacb55e158755eb79f77c41f8e/mypy-1.19.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c35d298c2c4bba75feb2195655dfea8124d855dfd7343bf8b8c055421eaf0cf8", size = 13620847, upload-time = "2025-12-15T05:03:39.633Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/42/332951aae42b79329f743bf1da088cd75d8d4d9acc18fbcbd84f26c1af4e/mypy-1.19.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:34c81968774648ab5ac09c29a375fdede03ba253f8f8287847bd480782f73a6a", size = 13834976, upload-time = "2025-12-15T05:03:08.786Z" },
+ { url = "https://files.pythonhosted.org/packages/6f/63/e7493e5f90e1e085c562bb06e2eb32cae27c5057b9653348d38b47daaecc/mypy-1.19.1-cp312-cp312-win_amd64.whl", hash = "sha256:b10e7c2cd7870ba4ad9b2d8a6102eb5ffc1f16ca35e3de6bfa390c1113029d13", size = 10118104, upload-time = "2025-12-15T05:03:10.834Z" },
+ { url = "https://files.pythonhosted.org/packages/de/9f/a6abae693f7a0c697dbb435aac52e958dc8da44e92e08ba88d2e42326176/mypy-1.19.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e3157c7594ff2ef1634ee058aafc56a82db665c9438fd41b390f3bde1ab12250", size = 13201927, upload-time = "2025-12-15T05:02:29.138Z" },
+ { url = "https://files.pythonhosted.org/packages/9a/a4/45c35ccf6e1c65afc23a069f50e2c66f46bd3798cbe0d680c12d12935caa/mypy-1.19.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bdb12f69bcc02700c2b47e070238f42cb87f18c0bc1fc4cdb4fb2bc5fd7a3b8b", size = 12206730, upload-time = "2025-12-15T05:03:01.325Z" },
+ { url = "https://files.pythonhosted.org/packages/05/bb/cdcf89678e26b187650512620eec8368fded4cfd99cfcb431e4cdfd19dec/mypy-1.19.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f859fb09d9583a985be9a493d5cfc5515b56b08f7447759a0c5deaf68d80506e", size = 12724581, upload-time = "2025-12-15T05:03:20.087Z" },
+ { url = "https://files.pythonhosted.org/packages/d1/32/dd260d52babf67bad8e6770f8e1102021877ce0edea106e72df5626bb0ec/mypy-1.19.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9a6538e0415310aad77cb94004ca6482330fece18036b5f360b62c45814c4ef", size = 13616252, upload-time = "2025-12-15T05:02:49.036Z" },
+ { url = "https://files.pythonhosted.org/packages/71/d0/5e60a9d2e3bd48432ae2b454b7ef2b62a960ab51292b1eda2a95edd78198/mypy-1.19.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:da4869fc5e7f62a88f3fe0b5c919d1d9f7ea3cef92d3689de2823fd27e40aa75", size = 13840848, upload-time = "2025-12-15T05:02:55.95Z" },
+ { url = "https://files.pythonhosted.org/packages/98/76/d32051fa65ecf6cc8c6610956473abdc9b4c43301107476ac03559507843/mypy-1.19.1-cp313-cp313-win_amd64.whl", hash = "sha256:016f2246209095e8eda7538944daa1d60e1e8134d98983b9fc1e92c1fc0cb8dd", size = 10135510, upload-time = "2025-12-15T05:02:58.438Z" },
+ { url = "https://files.pythonhosted.org/packages/de/eb/b83e75f4c820c4247a58580ef86fcd35165028f191e7e1ba57128c52782d/mypy-1.19.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06e6170bd5836770e8104c8fdd58e5e725cfeb309f0a6c681a811f557e97eac1", size = 13199744, upload-time = "2025-12-15T05:03:30.823Z" },
+ { url = "https://files.pythonhosted.org/packages/94/28/52785ab7bfa165f87fcbb61547a93f98bb20e7f82f90f165a1f69bce7b3d/mypy-1.19.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:804bd67b8054a85447c8954215a906d6eff9cabeabe493fb6334b24f4bfff718", size = 12215815, upload-time = "2025-12-15T05:02:42.323Z" },
+ { url = "https://files.pythonhosted.org/packages/0a/c6/bdd60774a0dbfb05122e3e925f2e9e846c009e479dcec4821dad881f5b52/mypy-1.19.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21761006a7f497cb0d4de3d8ef4ca70532256688b0523eee02baf9eec895e27b", size = 12740047, upload-time = "2025-12-15T05:03:33.168Z" },
+ { url = "https://files.pythonhosted.org/packages/32/2a/66ba933fe6c76bd40d1fe916a83f04fed253152f451a877520b3c4a5e41e/mypy-1.19.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:28902ee51f12e0f19e1e16fbe2f8f06b6637f482c459dd393efddd0ec7f82045", size = 13601998, upload-time = "2025-12-15T05:03:13.056Z" },
+ { url = "https://files.pythonhosted.org/packages/e3/da/5055c63e377c5c2418760411fd6a63ee2b96cf95397259038756c042574f/mypy-1.19.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:481daf36a4c443332e2ae9c137dfee878fcea781a2e3f895d54bd3002a900957", size = 13807476, upload-time = "2025-12-15T05:03:17.977Z" },
+ { url = "https://files.pythonhosted.org/packages/cd/09/4ebd873390a063176f06b0dbf1f7783dd87bd120eae7727fa4ae4179b685/mypy-1.19.1-cp314-cp314-win_amd64.whl", hash = "sha256:8bb5c6f6d043655e055be9b542aa5f3bdd30e4f3589163e85f93f3640060509f", size = 10281872, upload-time = "2025-12-15T05:03:05.549Z" },
+ { url = "https://files.pythonhosted.org/packages/8d/f4/4ce9a05ce5ded1de3ec1c1d96cf9f9504a04e54ce0ed55cfa38619a32b8d/mypy-1.19.1-py3-none-any.whl", hash = "sha256:f1235f5ea01b7db5468d53ece6aaddf1ad0b88d9e7462b86ef96fe04995d7247", size = 2471239, upload-time = "2025-12-15T05:03:07.248Z" },
+]
+
+[[package]]
+name = "mypy-extensions"
+version = "1.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" },
+]
+
[[package]]
name = "nodeenv"
version = "1.10.0"
@@ -1821,6 +1961,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" },
]
+[[package]]
+name = "nodejs-wheel-binaries"
+version = "24.13.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e5/d0/81d98b8fddc45332f79d6ad5749b1c7409fb18723545eae75d9b7e0048fb/nodejs_wheel_binaries-24.13.1.tar.gz", hash = "sha256:512659a67449a038231e2e972d49e77049d2cf789ae27db39eff4ab1ca52ac57", size = 8056, upload-time = "2026-02-12T17:31:04.368Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/aa/04/1ffe1838306654fcb50bcf46172567d50c8e27a76f4b9e55a1971fab5c4f/nodejs_wheel_binaries-24.13.1-py2.py3-none-macosx_13_0_arm64.whl", hash = "sha256:360ac9382c651de294c23c4933a02358c4e11331294983f3cf50ca1ac32666b1", size = 54757440, upload-time = "2026-02-12T17:30:35.748Z" },
+ { url = "https://files.pythonhosted.org/packages/66/f6/81ad81bc3bd919a20b110130c4fd318c7b6a5abb37eb53daa353ad908012/nodejs_wheel_binaries-24.13.1-py2.py3-none-macosx_13_0_x86_64.whl", hash = "sha256:035b718946793986762cdd50deee7f5f1a8f1b0bad0f0cfd57cad5492f5ea018", size = 54932957, upload-time = "2026-02-12T17:30:40.114Z" },
+ { url = "https://files.pythonhosted.org/packages/14/be/8e8a2bd50953c4c5b7e0fca07368d287917b84054dc3c93dd26a2940f0f9/nodejs_wheel_binaries-24.13.1-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:f795e9238438c4225f76fbd01e2b8e1a322116bbd0dc15a7dbd585a3ad97961e", size = 59287257, upload-time = "2026-02-12T17:30:43.781Z" },
+ { url = "https://files.pythonhosted.org/packages/58/57/92f6dfa40647702a9fa6d32393ce4595d0fc03c1daa9b245df66cc60e959/nodejs_wheel_binaries-24.13.1-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:978328e3ad522571eb163b042dfbd7518187a13968fe372738f90fdfe8a46afc", size = 59781783, upload-time = "2026-02-12T17:30:47.387Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/a5/457b984cf675cf86ace7903204b9c36edf7a2d1b4325ddf71eaf8d1027c7/nodejs_wheel_binaries-24.13.1-py2.py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e1dc893df85299420cd2a5feea0c3f8482a719b5f7f82d5977d58718b8b78b5f", size = 61287166, upload-time = "2026-02-12T17:30:50.646Z" },
+ { url = "https://files.pythonhosted.org/packages/3c/99/da515f7bc3bce35cfa6005f0e0c4e3c4042a466782b143112eb393b663be/nodejs_wheel_binaries-24.13.1-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:0e581ae219a39073dcadd398a2eb648f0707b0f5d68c565586139f919c91cbe9", size = 61870142, upload-time = "2026-02-12T17:30:54.563Z" },
+ { url = "https://files.pythonhosted.org/packages/cc/c0/22001d2c96d8200834af7d1de5e72daa3266c7270330275104c3d9ddd143/nodejs_wheel_binaries-24.13.1-py2.py3-none-win_amd64.whl", hash = "sha256:d4c969ea0bcb8c8b20bc6a7b4ad2796146d820278f17d4dc20229b088c833e22", size = 41185473, upload-time = "2026-02-12T17:30:57.524Z" },
+ { url = "https://files.pythonhosted.org/packages/ab/c4/7532325f968ecfc078e8a028e69a52e4c3f95fb800906bf6931ac1e89e2b/nodejs_wheel_binaries-24.13.1-py2.py3-none-win_arm64.whl", hash = "sha256:caec398cb9e94c560bacdcba56b3828df22a355749eb291f47431af88cbf26dc", size = 38881194, upload-time = "2026-02-12T17:31:00.214Z" },
+]
+
[[package]]
name = "openai"
version = "2.13.0"
@@ -1881,12 +2037,19 @@ dependencies = [
[package.optional-dependencies]
dev = [
{ name = "anyio" },
+ { name = "fastapi" },
{ name = "fastmcp" },
+ { name = "mypy" },
{ name = "pre-commit" },
+ { name = "pyarrow" },
+ { name = "pydantic-settings" },
+ { name = "pyright", extra = ["nodejs"] },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-cov" },
+ { name = "rich" },
{ name = "ruff" },
+ { name = "uvicorn" },
]
full = [
{ name = "fastapi" },
@@ -1912,17 +2075,19 @@ requires-dist = [
{ name = "anyio", marker = "extra == 'dev'", specifier = ">=4.0.0" },
{ name = "fastapi", marker = "extra == 'full'", specifier = ">=0.100.0,<1.0.0" },
{ name = "fastapi", marker = "extra == 'server'", specifier = ">=0.100.0,<1.0.0" },
- { name = "fastmcp", marker = "extra == 'dev'", specifier = ">=2.0.0" },
{ name = "fastmcp", marker = "extra == 'full'", specifier = ">=2.0.0" },
{ name = "fastmcp", marker = "extra == 'mcp'", specifier = ">=2.0.0" },
{ name = "httpx", specifier = ">=0.25.0,<1.0.0" },
{ name = "litellm", specifier = ">=1.40.0,<2.0.0" },
+ { name = "mypy", marker = "extra == 'dev'", specifier = "==1.19.1" },
+ { name = "osmosis-ai", extras = ["mcp", "server"], marker = "extra == 'dev'" },
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=4.0.0,<5.0.0" },
{ name = "pyarrow", marker = "extra == 'full'", specifier = ">=14.0.0" },
{ name = "pyarrow", marker = "extra == 'server'", specifier = ">=14.0.0" },
{ name = "pydantic", specifier = ">=2.0.0,<3.0.0" },
{ name = "pydantic-settings", marker = "extra == 'full'", specifier = ">=2.0.0,<3.0.0" },
{ name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.0.0,<3.0.0" },
+ { name = "pyright", extras = ["nodejs"], marker = "extra == 'dev'", specifier = "==1.1.408" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0,<10.0.0" },
{ name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0,<2.0.0" },
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=6.0.0" },
@@ -1956,6 +2121,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/7d/eb/b6260b31b1a96386c0a880edebe26f89669098acea8e0318bff6adb378fd/pathable-0.4.4-py3-none-any.whl", hash = "sha256:5ae9e94793b6ef5a4cbe0a7ce9dbbefc1eec38df253763fd0aeeacf2762dbbc2", size = 9592, upload-time = "2025-01-10T18:43:11.88Z" },
]
+[[package]]
+name = "pathspec"
+version = "1.0.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fa/36/e27608899f9b8d4dff0617b2d9ab17ca5608956ca44461ac14ac48b44015/pathspec-1.0.4.tar.gz", hash = "sha256:0210e2ae8a21a9137c0d470578cb0e595af87edaa6ebf12ff176f14a02e0e645", size = 131200, upload-time = "2026-01-27T03:59:46.938Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl", hash = "sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723", size = 55206, upload-time = "2026-01-27T03:59:45.137Z" },
+]
+
[[package]]
name = "pathvalidate"
version = "3.3.1"
@@ -2437,6 +2611,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063, upload-time = "2025-09-26T14:40:36.069Z" },
]
+[[package]]
+name = "pyright"
+version = "1.1.408"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "nodeenv" },
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/74/b2/5db700e52554b8f025faa9c3c624c59f1f6c8841ba81ab97641b54322f16/pyright-1.1.408.tar.gz", hash = "sha256:f28f2321f96852fa50b5829ea492f6adb0e6954568d1caa3f3af3a5f555eb684", size = 4400578, upload-time = "2026-01-08T08:07:38.795Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/0c/82/a2c93e32800940d9573fb28c346772a14778b84ba7524e691b324620ab89/pyright-1.1.408-py3-none-any.whl", hash = "sha256:090b32865f4fdb1e0e6cd82bf5618480d48eecd2eb2e70f960982a3d9a4c17c1", size = 6399144, upload-time = "2026-01-08T08:07:37.082Z" },
+]
+
+[package.optional-dependencies]
+nodejs = [
+ { name = "nodejs-wheel-binaries" },
+]
+
[[package]]
name = "pytest"
version = "9.0.2"
From 7b92f4d2ce74d600e539c28f01b61b8c2f27434e Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 00:15:13 -0800
Subject: [PATCH 25/60] Update CONTRIBUTING.md to streamline setup and testing
instructions
- Renamed "Setup" to "Quick Start" for clarity.
- Consolidated installation instructions for both `uv` and `pip`.
- Enhanced the commands reference section for better usability.
- Clarified type checking requirements for Pyright and mypy.
- Updated pre-commit hook instructions for improved clarity.
---
.github/dependabot.yml | 33 ++++++++++++++
CONTRIBUTING.md | 99 +++++++++++++++++-------------------------
2 files changed, 72 insertions(+), 60 deletions(-)
create mode 100644 .github/dependabot.yml
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 00000000..0f717257
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,33 @@
+version: 2
+updates:
+ # Python dependencies (pyproject.toml)
+ - package-ecosystem: pip
+ directory: "/"
+ schedule:
+ interval: weekly
+ day: monday
+ groups:
+ dev-dependencies:
+ patterns:
+ - "ruff"
+ - "pytest*"
+ - "mypy"
+ - "pyright*"
+ - "pre-commit"
+ commit-message:
+ prefix: "chore(deps):"
+ labels:
+ - "dependencies"
+ open-pull-requests-limit: 10
+
+ # GitHub Actions
+ - package-ecosystem: github-actions
+ directory: "/"
+ schedule:
+ interval: weekly
+ day: monday
+ commit-message:
+ prefix: "chore(ci):"
+ labels:
+ - "dependencies"
+ - "ci"
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 75cd08ea..2bb23b8b 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,100 +1,79 @@
# Contributing
-## Setup
+## Quick Start
+
+**Using uv (recommended):**
```bash
git clone https://github.com/Osmosis-AI/osmosis-sdk-python
cd osmosis-sdk-python
-
-# Install with all development dependencies (includes server + mcp extras
-# for type checking and full test coverage)
-pip install -e ".[dev]"
-
-# Or using uv (recommended)
uv sync --extra dev
-
-# Install pre-commit hooks
pre-commit install
+uv run pytest
```
-## Testing
+**Using pip:**
```bash
-# Run all tests
-pytest tests/
-
-# Run a single test file
-pytest tests/unit/rollout/core/test_base.py
-
-# Run tests matching a pattern
-pytest -k "test_name"
-
-# Run with coverage report
-pytest --cov=osmosis_ai --cov-report=term-missing
+git clone https://github.com/Osmosis-AI/osmosis-sdk-python
+cd osmosis-sdk-python
+python -m venv .venv && source .venv/bin/activate
+pip install -e ".[dev]"
+pre-commit install
+pytest
```
-Coverage configuration is in `pyproject.toml` under `[tool.coverage.*]`. CI enforces a minimum coverage threshold.
+## Commands Reference
-## Linting & Formatting
+The table below lists all development commands. If you installed with **pip**, drop the `uv run` prefix.
-This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and code formatting. Configuration lives in `pyproject.toml` under `[tool.ruff]`.
+| Task | Command |
+|------|---------|
+| Run all tests | `uv run pytest` |
+| Run a single file | `uv run pytest tests/unit/rollout/core/test_base.py` |
+| Run tests by name | `uv run pytest -k "test_name"` |
+| Run with coverage | `uv run pytest --cov=osmosis_ai --cov-report=term-missing` |
+| Lint | `uv run ruff check .` |
+| Lint + autofix | `uv run ruff check --fix .` |
+| Format | `uv run ruff format .` |
+| Check formatting | `uv run ruff format --check .` |
+| Type check (pyright) | `uv run pyright osmosis_ai/` |
+| Type check (mypy) | `uv run mypy osmosis_ai/` |
-```bash
-# Check for lint errors
-ruff check .
+## Testing
-# Auto-fix lint errors where possible
-ruff check --fix .
+Coverage configuration is in `pyproject.toml` under `[tool.coverage.*]`. CI enforces a minimum coverage threshold of 70%.
-# Format code
-ruff format .
+## Linting & Formatting
-# Check formatting without modifying files
-ruff format --check .
-```
+This project uses [Ruff](https://docs.astral.sh/ruff/) for both linting and code formatting. Configuration lives in `pyproject.toml` under `[tool.ruff]`.
Ruff is pinned to one version across `pyproject.toml`, `.pre-commit-config.yaml`, and CI so local checks, pre-commit hooks, and GitHub Actions produce the same results.
## Type Checking
-This project uses [Pyright](https://microsoft.github.io/pyright/) as the primary type checker and [mypy](https://mypy-lang.org/) as a secondary checker. Both are included in the `dev` extras — install with:
-
-```bash
-uv sync --extra dev
-```
-
-Then run via `uv run`:
+[Pyright](https://microsoft.github.io/pyright/) is the primary type checker and [mypy](https://mypy-lang.org/) is a secondary checker. Both are included in the `dev` extras.
-```bash
-# Pyright — must pass (blocking in CI)
-uv run pyright osmosis_ai/
-
-# mypy — advisory (non-blocking in CI)
-uv run mypy osmosis_ai/
-```
-
-Pyright runs in `standard` mode. All pyright errors must be resolved before merging. mypy runs in CI with `continue-on-error` and serves as an advisory check — fix mypy warnings when practical, but they won't block a PR.
+- **Pyright** — must pass. All errors must be resolved before merging.
+- **mypy** — advisory (`continue-on-error` in CI). Fix warnings when practical, but they won't block a PR.
Configuration for both tools lives in `pyproject.toml` under `[tool.pyright]` and `[tool.mypy]`.
-> **Note:** CI also runs `pyright --verifytypes osmosis_ai --ignoreexternal` to check public API type completeness, but this command requires a non-editable install to locate the `py.typed` marker. It does not work with local editable installs (`uv sync`) and is therefore non-blocking in CI.
+> **Note:** CI also runs `pyright --verifytypes osmosis_ai --ignoreexternal` to check public API type completeness. This requires a non-editable install and is non-blocking in CI.
## Pre-commit Hooks
-A [pre-commit](https://pre-commit.com/) configuration is included to run Ruff automatically on every commit:
+[Pre-commit](https://pre-commit.com/) runs `ruff check --fix` and `ruff format` automatically on every commit. Make sure hooks are installed before submitting a pull request:
```bash
pre-commit install
```
-This ensures `ruff check --fix` and `ruff format` run before each commit. Please make sure hooks are installed before submitting a pull request.
-
## Pull Requests
-1. Fork the repository
-2. Create a feature branch
-3. Make your changes
-4. Run tests and linting
-5. Submit a pull request
+1. Fork the repository and create a feature branch
+2. Make your changes
+3. Run `uv run pytest` and `uv run ruff check .`
+4. Submit a pull request
-CI will run linting (`ruff check` + `ruff format --check`), type checking (pyright + mypy), and tests with coverage on every PR.
+CI will run linting, type checking (pyright + mypy), tests across Python 3.10–3.13, and a build validation on every PR.
From 120bf30ce4475cb1a6f41f578daa823e8fefe703 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 10:13:55 -0800
Subject: [PATCH 26/60] Enhance type annotations and add type stubs for PyYAML
and requests
- Added type stubs for `PyYAML` and `requests` in `pyproject.toml` and `uv.lock` to improve type checking in development.
- Updated various files to include type hints for better clarity and type safety, including changes in `cli.py`, `platform_client.py`, and several other modules.
- Refactored return types and variable annotations to ensure consistency across the codebase.
---
osmosis_ai/auth/platform_client.py | 3 ++-
osmosis_ai/cli.py | 3 ++-
osmosis_ai/cli_services/config.py | 6 ++---
osmosis_ai/cli_services/engine.py | 10 +++++---
osmosis_ai/rollout/client.py | 9 ++++---
osmosis_ai/rollout/console.py | 2 ++
osmosis_ai/rollout/eval/common/dataset.py | 6 +++--
osmosis_ai/rollout/eval/evaluation/eval_fn.py | 4 ++-
osmosis_ai/rollout/eval/evaluation/runner.py | 4 ++-
.../rollout/eval/test_mode/interactive.py | 3 ++-
osmosis_ai/rollout/mcp/agent_loop.py | 3 ++-
osmosis_ai/rollout/mcp/loader.py | 3 ++-
osmosis_ai/rollout/server/app.py | 4 +--
osmosis_ai/rollout/testing.py | 2 +-
pyproject.toml | 2 ++
uv.lock | 25 +++++++++++++++++++
16 files changed, 67 insertions(+), 22 deletions(-)
diff --git a/osmosis_ai/auth/platform_client.py b/osmosis_ai/auth/platform_client.py
index 5a5c6846..0ca0587c 100644
--- a/osmosis_ai/auth/platform_client.py
+++ b/osmosis_ai/auth/platform_client.py
@@ -95,7 +95,8 @@ def platform_request(
try:
with urlopen(request, timeout=timeout) as response:
- return json.loads(response.read().decode())
+ result: dict[str, Any] = json.loads(response.read().decode())
+ return result
except HTTPError as e:
if e.code == 401:
_handle_401_and_cleanup()
diff --git a/osmosis_ai/cli.py b/osmosis_ai/cli.py
index 23551e19..4c922183 100644
--- a/osmosis_ai/cli.py
+++ b/osmosis_ai/cli.py
@@ -31,7 +31,8 @@ def main(argv: list[str] | None = None) -> int:
return 1
try:
- return handler(args)
+ exit_code: int = handler(args)
+ return exit_code
except CLIError as exc:
print(f"Error: {exc}", file=sys.stderr)
return 1
diff --git a/osmosis_ai/cli_services/config.py b/osmosis_ai/cli_services/config.py
index e0159780..099921e4 100644
--- a/osmosis_ai/cli_services/config.py
+++ b/osmosis_ai/cli_services/config.py
@@ -76,7 +76,7 @@ def parse_document(
class BaseRubricConfigSchema(RubricConfigDocumentSchema):
"""Schema handling documents without an explicit version."""
- version = None
+ version: int | None = None
def parse_document(
self,
@@ -96,7 +96,7 @@ def parse_document(
class Version1RubricConfigSchema(BaseRubricConfigSchema):
"""Schema for version 1 documents."""
- version = 1
+ version: int | None = 1
class RubricConfigParser:
@@ -423,7 +423,7 @@ class _LiteralSafeDumper(yaml.SafeDumper):
"""YAML dumper that preserves multiline strings with literal blocks."""
-def _represent_str(dumper: _LiteralSafeDumper, data: str): # type: ignore[type-arg]
+def _represent_str(dumper: _LiteralSafeDumper, data: str) -> Any:
if "\n" in data:
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
return SafeRepresenter.represent_str(dumper, data)
diff --git a/osmosis_ai/cli_services/engine.py b/osmosis_ai/cli_services/engine.py
index eb834f1c..24887155 100644
--- a/osmosis_ai/cli_services/engine.py
+++ b/osmosis_ai/cli_services/engine.py
@@ -203,7 +203,7 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
)
if is_evaluate_rubric_style:
- return self._evaluate_fn(
+ result: dict[str, Any] = self._evaluate_fn(
rubric=config.rubric_text,
solution_str=solution,
model_info=model_info_payload,
@@ -214,6 +214,7 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
score_max=score_max,
return_details=True,
)
+ return result
call_args: list[Any] = []
call_kwargs: dict[str, Any] = {}
@@ -225,7 +226,7 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
continue
if param.name == "solution_str":
- value = solution
+ value: Any = solution
elif param.name == "ground_truth":
value = ground_truth
elif param.name == "extra_info":
@@ -241,7 +242,10 @@ def run(self, config: RubricConfig, record: DatasetRecord) -> dict[str, Any]:
else:
call_kwargs[param.name] = value
- return self._evaluate_fn(*call_args, **call_kwargs)
+ fallback_result: dict[str, Any] = self._evaluate_fn(
+ *call_args, **call_kwargs
+ )
+ return fallback_result
except (MissingAPIKeyError, ProviderRequestError, ModelNotFoundError) as exc:
raise CLIError(str(exc)) from exc
diff --git a/osmosis_ai/rollout/client.py b/osmosis_ai/rollout/client.py
index 9b6a1b58..e62173f2 100644
--- a/osmosis_ai/rollout/client.py
+++ b/osmosis_ai/rollout/client.py
@@ -87,7 +87,8 @@ def has_tool_calls(self) -> bool:
@property
def tool_calls(self) -> list[dict[str, Any]]:
"""Get tool calls from the response."""
- return self.message.get("tool_calls", [])
+ calls: list[dict[str, Any]] = self.message.get("tool_calls", [])
+ return calls
@property
def content(self) -> str | None:
@@ -173,8 +174,8 @@ def __init__(
# Settings for connection pool
self._max_connections = settings.max_connections
self._max_keepalive_connections = settings.max_keepalive_connections
- self._retry_base_delay = settings.retry_base_delay
- self._retry_max_delay = settings.retry_max_delay
+ self._retry_base_delay: float = settings.retry_base_delay
+ self._retry_max_delay: float = settings.retry_max_delay
# HTTP client (lazy initialized)
self._client: httpx.AsyncClient | None = None
@@ -218,7 +219,7 @@ async def close(self) -> None:
def _calculate_retry_delay(self, attempt: int) -> float:
"""Calculate delay for retry with exponential backoff."""
delay = self._retry_base_delay * (2**attempt)
- return min(delay, self._retry_max_delay)
+ return float(min(delay, self._retry_max_delay))
async def chat_completions(
self,
diff --git a/osmosis_ai/rollout/console.py b/osmosis_ai/rollout/console.py
index c46c5c16..31b1f240 100644
--- a/osmosis_ai/rollout/console.py
+++ b/osmosis_ai/rollout/console.py
@@ -105,6 +105,8 @@ def __init__(
# Use rich if available and writing to TTY
self._use_rich = RICH_AVAILABLE and self._is_tty and not no_color
+ self._rich: RichConsole | None
+ self._rich_stderr: RichConsole | None
if self._use_rich:
self._rich = RichConsole(file=self._file, force_terminal=force_terminal)
self._rich_stderr = RichConsole(file=sys.stderr)
diff --git a/osmosis_ai/rollout/eval/common/dataset.py b/osmosis_ai/rollout/eval/common/dataset.py
index 58b6718c..228fa06c 100644
--- a/osmosis_ai/rollout/eval/common/dataset.py
+++ b/osmosis_ai/rollout/eval/common/dataset.py
@@ -159,7 +159,8 @@ def _parse_parquet(self) -> list[dict[str, Any]]:
try:
table = pq.read_table(self.file_path)
- return table.to_pylist()
+ rows: list[dict[str, Any]] = table.to_pylist()
+ return rows
except Exception as e:
raise DatasetParseError(f"Error reading Parquet file: {e}") from e
@@ -173,7 +174,8 @@ def _count_parquet_rows(self) -> int:
try:
metadata = pq.read_metadata(self.file_path)
- return metadata.num_rows
+ count: int = metadata.num_rows
+ return count
except Exception as e:
raise DatasetParseError(f"Error reading Parquet metadata: {e}") from e
diff --git a/osmosis_ai/rollout/eval/evaluation/eval_fn.py b/osmosis_ai/rollout/eval/evaluation/eval_fn.py
index c3305198..78e1bfd9 100644
--- a/osmosis_ai/rollout/eval/evaluation/eval_fn.py
+++ b/osmosis_ai/rollout/eval/evaluation/eval_fn.py
@@ -94,6 +94,7 @@ async def __call__(
Returns:
Float score from the eval function.
"""
+ kwargs: dict[str, Any]
if self._mode == "simple":
solution_str = self._extract_last_assistant_content(messages)
kwargs = {
@@ -162,7 +163,8 @@ def load_eval_fn(module_path: str) -> Callable:
if not callable(fn):
raise EvalFnError(f"'{attr_name}' in '{module_name}' is not callable")
- return fn
+ result: Callable[..., Any] = fn
+ return result
def load_eval_fns(module_paths: list[str]) -> list[EvalFnWrapper]:
diff --git a/osmosis_ai/rollout/eval/evaluation/runner.py b/osmosis_ai/rollout/eval/evaluation/runner.py
index f0e9e52c..4c665384 100644
--- a/osmosis_ai/rollout/eval/evaluation/runner.py
+++ b/osmosis_ai/rollout/eval/evaluation/runner.py
@@ -378,7 +378,9 @@ async def run_eval(
total_start = time.monotonic()
row_results: list[EvalRowResult] = []
- model_tags = ["primary", "baseline"] if self.has_baseline else [None]
+ model_tags: list[str | None] = (
+ ["primary", "baseline"] if self.has_baseline else [None]
+ )
total = len(rows) * n_runs * len(model_tags)
current = 0
stopped_early = False
diff --git a/osmosis_ai/rollout/eval/test_mode/interactive.py b/osmosis_ai/rollout/eval/test_mode/interactive.py
index 03eb5eaf..4faae8fa 100644
--- a/osmosis_ai/rollout/eval/test_mode/interactive.py
+++ b/osmosis_ai/rollout/eval/test_mode/interactive.py
@@ -474,7 +474,8 @@ def _handle_step(self, step: InteractiveStep) -> bool:
handler = getattr(self, handler_name)
should_return, return_value = handler()
if should_return:
- return return_value
+ result: bool = return_value
+ return result
else:
self._print_help(user_input)
diff --git a/osmosis_ai/rollout/mcp/agent_loop.py b/osmosis_ai/rollout/mcp/agent_loop.py
index d1d16a95..74b636b8 100644
--- a/osmosis_ai/rollout/mcp/agent_loop.py
+++ b/osmosis_ai/rollout/mcp/agent_loop.py
@@ -39,7 +39,8 @@ def _resolve_ref(ref: str, defs: dict[str, Any]) -> dict[str, Any]:
node: Any = {"$defs": defs}
for part in parts:
node = node[part]
- return node
+ result: dict[str, Any] = node
+ return result
def _json_schema_to_property(
diff --git a/osmosis_ai/rollout/mcp/loader.py b/osmosis_ai/rollout/mcp/loader.py
index a72c19db..4ff6970b 100644
--- a/osmosis_ai/rollout/mcp/loader.py
+++ b/osmosis_ai/rollout/mcp/loader.py
@@ -12,6 +12,7 @@
import os
import sys
import types
+from collections.abc import Iterator
from typing import Any
@@ -49,7 +50,7 @@ def _clear_new_local_imports(
@contextlib.contextmanager
-def _temporary_sys_path(path: str):
+def _temporary_sys_path(path: str) -> Iterator[None]:
"""Temporarily prepend ``path`` to ``sys.path`` and restore exactly."""
original = list(sys.path)
# Ensure local sibling imports resolve to the MCP directory during import.
diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py
index a0bd38fe..29f7a2d0 100644
--- a/osmosis_ai/rollout/server/app.py
+++ b/osmosis_ai/rollout/server/app.py
@@ -17,7 +17,7 @@
import asyncio
import logging
import time
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any
@@ -151,7 +151,7 @@ def create_app(
)
@asynccontextmanager
- async def lifespan(app: FastAPI):
+ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
"""Manage application lifecycle."""
logger.info(
"Server starting: agent_loop=%s, max_concurrent=%d",
diff --git a/osmosis_ai/rollout/testing.py b/osmosis_ai/rollout/testing.py
index e24de716..2a9c992c 100644
--- a/osmosis_ai/rollout/testing.py
+++ b/osmosis_ai/rollout/testing.py
@@ -295,7 +295,7 @@ def mock_trainer(monkeypatch):
original_post = httpx.AsyncClient.post
- async def mock_post(self, url: str, **kwargs):
+ async def mock_post(self: Any, url: str, **kwargs: Any) -> Any:
if "/v1/chat/completions" in url:
resp = client.post("/v1/chat/completions", **kwargs)
return httpx.Response(
diff --git a/pyproject.toml b/pyproject.toml
index 59f1c577..44d92478 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -55,6 +55,8 @@ dev = [
"pre-commit>=4.0.0,<5.0.0",
"pyright[nodejs]==1.1.408",
"mypy==1.19.1",
+ "types-PyYAML>=6.0",
+ "types-requests>=2.0",
]
# Full installation with all optional features (union of server + mcp)
diff --git a/uv.lock b/uv.lock
index 0c632111..d86de8d2 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2049,6 +2049,8 @@ dev = [
{ name = "pytest-cov" },
{ name = "rich" },
{ name = "ruff" },
+ { name = "types-pyyaml" },
+ { name = "types-requests" },
{ name = "uvicorn" },
]
full = [
@@ -2098,6 +2100,8 @@ requires-dist = [
{ name = "rich", marker = "extra == 'server'", specifier = ">=13.0.0" },
{ name = "ruff", marker = "extra == 'dev'", specifier = "==0.15.1" },
{ name = "tqdm", specifier = ">=4.0.0,<5.0.0" },
+ { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0" },
+ { name = "types-requests", marker = "extra == 'dev'", specifier = ">=2.0" },
{ name = "uvicorn", marker = "extra == 'full'", specifier = ">=0.23.0,<1.0.0" },
{ name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.23.0,<1.0.0" },
]
@@ -3407,6 +3411,27 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/0a/4aca634faf693e33004796b6cee0ae2e1dba375a800c16ab8d3eff4bb800/typer_slim-0.21.1-py3-none-any.whl", hash = "sha256:6e6c31047f171ac93cc5a973c9e617dbc5ab2bddc4d0a3135dc161b4e2020e0d", size = 47444, upload-time = "2026-01-06T11:21:12.441Z" },
]
+[[package]]
+name = "types-pyyaml"
+version = "6.0.12.20250915"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7e/69/3c51b36d04da19b92f9e815be12753125bd8bc247ba0470a982e6979e71c/types_pyyaml-6.0.12.20250915.tar.gz", hash = "sha256:0f8b54a528c303f0e6f7165687dd33fafa81c807fcac23f632b63aa624ced1d3", size = 17522, upload-time = "2025-09-15T03:01:00.728Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl", hash = "sha256:e7d4d9e064e89a3b3cae120b4990cd370874d2bf12fa5f46c97018dd5d3c9ab6", size = 20338, upload-time = "2025-09-15T03:00:59.218Z" },
+]
+
+[[package]]
+name = "types-requests"
+version = "2.32.4.20260107"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/0f/f3/a0663907082280664d745929205a89d41dffb29e89a50f753af7d57d0a96/types_requests-2.32.4.20260107.tar.gz", hash = "sha256:018a11ac158f801bfa84857ddec1650750e393df8a004a8a9ae2a9bec6fcb24f", size = 23165, upload-time = "2026-01-07T03:20:54.091Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/1c/12/709ea261f2bf91ef0a26a9eed20f2623227a8ed85610c1e54c5805692ecb/types_requests-2.32.4.20260107-py3-none-any.whl", hash = "sha256:b703fe72f8ce5b31ef031264fe9395cac8f46a04661a79f7ed31a80fb308730d", size = 20676, upload-time = "2026-01-07T03:20:52.929Z" },
+]
+
[[package]]
name = "typing-extensions"
version = "4.15.0"
From 02b5503b46a25fb18d3d610938655c7a23c18545 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 10:38:23 -0800
Subject: [PATCH 27/60] Refactor GitHub Actions workflow to streamline testing
- Removed the mypy type checking job from the CI workflow.
- Renamed the test job to 'pytest' and updated the Python version matrix to exclude 3.11.
- Enhanced the overall structure for better clarity and maintainability.
---
.github/workflows/tests.yml | 18 ++----------------
1 file changed, 2 insertions(+), 16 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 19797ebb..238abbed 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -38,26 +38,12 @@ jobs:
continue-on-error: true
run: uv run pyright --verifytypes osmosis_ai --ignoreexternal
- typecheck-mypy:
- runs-on: ubuntu-latest
- continue-on-error: true
- steps:
- - uses: actions/checkout@v4
- - uses: astral-sh/setup-uv@v7
- with:
- python-version: "3.12"
- enable-cache: true
- - name: Install dependencies
- run: uv sync --frozen --extra dev
- - name: Run mypy
- run: uv run mypy osmosis_ai/
-
- test:
+ pytest:
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
- python-version: ["3.10", "3.11", "3.12", "3.13"]
+ python-version: ["3.10", "3.12", "3.13"]
steps:
- uses: actions/checkout@v4
From 9023a1b8a218d652b8288b0b05e05738c452a82e Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 17 Feb 2026 18:45:29 +0000
Subject: [PATCH 28/60] chore(ci): bump actions/setup-python from 5 to 6
Bumps [actions/setup-python](https://github.com/actions/setup-python) from 5 to 6.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/v5...v6)
---
updated-dependencies:
- dependency-name: actions/setup-python
dependency-version: '6'
dependency-type: direct:production
update-type: version-update:semver-major
...
Signed-off-by: dependabot[bot]
---
.github/workflows/publish.yml | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 20e24228..8f3b7c6f 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -18,7 +18,7 @@ jobs:
uses: actions/checkout@v4
- name: Set up Python
- uses: actions/setup-python@v5
+ uses: actions/setup-python@v6
with:
python-version: "3.12"
From 44224c5a97f7654e98086801e40a5ca4c72ba817 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 17 Feb 2026 18:45:34 +0000
Subject: [PATCH 29/60] chore(ci): bump actions/checkout from 4 to 6
Bumps [actions/checkout](https://github.com/actions/checkout) from 4 to 6.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/v4...v6)
---
updated-dependencies:
- dependency-name: actions/checkout
dependency-version: '6'
dependency-type: direct:production
update-type: version-update:semver-major
...
Signed-off-by: dependabot[bot]
---
.github/workflows/publish.yml | 2 +-
.github/workflows/tests.yml | 8 ++++----
2 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 20e24228..3e2cbcfb 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -15,7 +15,7 @@ jobs:
steps:
- name: Checkout repository
- uses: actions/checkout@v4
+ uses: actions/checkout@v6
- name: Set up Python
uses: actions/setup-python@v5
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 238abbed..b5b30c56 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -10,7 +10,7 @@ jobs:
lint:
runs-on: ubuntu-latest
steps:
- - uses: actions/checkout@v4
+ - uses: actions/checkout@v6
- uses: astral-sh/ruff-action@v3
@@ -21,7 +21,7 @@ jobs:
typecheck-pyright:
runs-on: ubuntu-latest
steps:
- - uses: actions/checkout@v4
+ - uses: actions/checkout@v6
- uses: astral-sh/setup-uv@v7
with:
python-version: "3.12"
@@ -45,7 +45,7 @@ jobs:
matrix:
python-version: ["3.10", "3.12", "3.13"]
steps:
- - uses: actions/checkout@v4
+ - uses: actions/checkout@v6
- uses: astral-sh/setup-uv@v7
with:
@@ -66,7 +66,7 @@ jobs:
build:
runs-on: ubuntu-latest
steps:
- - uses: actions/checkout@v4
+ - uses: actions/checkout@v6
- uses: astral-sh/setup-uv@v7
with:
From 29dcdf94a770358b71b5aa60eab085ee6161d3bb Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 11:00:21 -0800
Subject: [PATCH 30/60] Update GitHub Actions workflow to improve testing
efficiency
- Reintroduced the mypy type checking job in the CI workflow.
- Adjusted the Python version matrix to include 3.11 and removed 3.12.
- Refined the structure of the workflow for enhanced clarity and maintainability.
---
.github/CODEOWNERS | 13 +++
.github/ISSUE_TEMPLATE/bug_report.yml | 93 ++++++++++++++++++++++
.github/ISSUE_TEMPLATE/config.yml | 5 ++
.github/ISSUE_TEMPLATE/feature_request.yml | 56 +++++++++++++
.github/PULL_REQUEST_TEMPLATE.md | 29 +++++++
SECURITY.md | 33 ++++++++
6 files changed, 229 insertions(+)
create mode 100644 .github/CODEOWNERS
create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yml
create mode 100644 .github/ISSUE_TEMPLATE/config.yml
create mode 100644 .github/ISSUE_TEMPLATE/feature_request.yml
create mode 100644 .github/PULL_REQUEST_TEMPLATE.md
create mode 100644 SECURITY.md
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
new file mode 100644
index 00000000..cc3ad7c6
--- /dev/null
+++ b/.github/CODEOWNERS
@@ -0,0 +1,13 @@
+# Default: core maintainers review everything
+* @BaiqingL @JoyboyBrian @JakeTrock @allenlinsh @mathewjhan @pandyamarut @artem-osmosis
+
+# Rollout SDK (core training integration)
+/osmosis_ai/rollout/ @BaiqingL @JoyboyBrian @JakeTrock @allenlinsh @mathewjhan @pandyamarut @artem-osmosis
+
+# CI/CD and build configuration
+/.github/ @BaiqingL @JoyboyBrian @JakeTrock
+/pyproject.toml @BaiqingL @JoyboyBrian @JakeTrock
+
+# Documentation
+/README.md @BaiqingL @JoyboyBrian @JakeTrock
+/CONTRIBUTING.md @BaiqingL @JoyboyBrian @JakeTrock
diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml
new file mode 100644
index 00000000..27f600dd
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.yml
@@ -0,0 +1,93 @@
+name: Bug Report
+description: Report a bug in osmosis-ai SDK
+labels: ["bug"]
+body:
+ - type: markdown
+ attributes:
+ value: |
+ Thanks for reporting a bug! Please fill in the details below.
+
+ - type: textarea
+ id: description
+ attributes:
+ label: Description
+ description: A clear and concise description of the bug.
+ validations:
+ required: true
+
+ - type: textarea
+ id: reproduction
+ attributes:
+ label: Steps to Reproduce
+ description: Minimal steps to reproduce the behavior.
+ placeholder: |
+ 1. Install osmosis-ai with `pip install osmosis-ai`
+ 2. Run the following code:
+ ```python
+ from osmosis_ai import ...
+ ```
+ 3. See error
+ validations:
+ required: true
+
+ - type: textarea
+ id: expected
+ attributes:
+ label: Expected Behavior
+ description: What you expected to happen.
+ validations:
+ required: true
+
+ - type: textarea
+ id: actual
+ attributes:
+ label: Actual Behavior
+ description: What actually happened. Include full error messages or tracebacks if applicable.
+ validations:
+ required: true
+
+ - type: input
+ id: sdk-version
+ attributes:
+ label: SDK Version
+ description: "Output of `pip show osmosis-ai | grep Version` or `osmosis --version`"
+ placeholder: "e.g. 0.5.0"
+ validations:
+ required: true
+
+ - type: input
+ id: python-version
+ attributes:
+ label: Python Version
+ description: "Output of `python --version`"
+ placeholder: "e.g. 3.12.3"
+ validations:
+ required: true
+
+ - type: input
+ id: os
+ attributes:
+ label: OS
+ placeholder: "e.g. Ubuntu 22.04 / macOS 14.5 / Windows 11"
+ validations:
+ required: false
+
+ - type: dropdown
+ id: install-method
+ attributes:
+ label: Installation Method
+ options:
+ - pip install osmosis-ai
+ - pip install -e . (editable)
+ - uv sync
+ - Other
+ validations:
+ required: false
+
+ - type: textarea
+ id: context
+ attributes:
+ label: Additional Context
+ description: Any other context, logs, or screenshots.
+ validations:
+ required: false
diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml
new file mode 100644
index 00000000..7e37756e
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -0,0 +1,5 @@
+blank_issues_enabled: false
+contact_links:
+ - name: Documentation
+ url: https://github.com/Osmosis-AI/osmosis-sdk-python#readme
+ about: Check the README and docs before opening an issue.
diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml
new file mode 100644
index 00000000..371ea73c
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature_request.yml
@@ -0,0 +1,56 @@
+name: Feature Request
+description: Suggest a new feature or improvement
+labels: ["enhancement"]
+body:
+ - type: markdown
+ attributes:
+ value: |
+ Thanks for suggesting a feature! Please describe what you'd like.
+
+ - type: textarea
+ id: problem
+ attributes:
+ label: Problem or Use Case
+ description: What problem does this solve? What are you trying to accomplish?
+ placeholder: "I'm trying to ... but currently ..."
+ validations:
+ required: true
+
+ - type: textarea
+ id: solution
+ attributes:
+ label: Proposed Solution
+ description: Describe the feature or behavior you'd like to see.
+ validations:
+ required: true
+
+ - type: textarea
+ id: alternatives
+ attributes:
+ label: Alternatives Considered
+ description: Any alternative solutions or workarounds you've considered.
+ validations:
+ required: false
+
+ - type: dropdown
+ id: component
+ attributes:
+ label: SDK Component
+ description: Which part of the SDK does this relate to?
+ options:
+ - Reward / Rubric decorators
+ - Remote Rollout (RolloutAgentLoop)
+ - Rollout Server (create_app)
+ - CLI (osmosis command)
+ - Authentication
+ - Other
+ validations:
+ required: false
+
+ - type: textarea
+ id: context
+ attributes:
+ label: Additional Context
+ description: Any other context, code examples, or references.
+ validations:
+ required: false
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
new file mode 100644
index 00000000..ce6a50e2
--- /dev/null
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -0,0 +1,29 @@
+## What
+
+
+
+## Why
+
+
+
+
+## Type of Change
+
+- [ ] Bug fix (non-breaking change that fixes an issue)
+- [ ] New feature (non-breaking change that adds functionality)
+- [ ] Breaking change (fix or feature that would cause existing functionality to change)
+- [ ] Refactor (code change that neither fixes a bug nor adds a feature)
+- [ ] Documentation update
+- [ ] CI/build configuration change
+
+## How to Test
+
+
+
+## Checklist
+
+- [ ] `ruff check .` and `ruff format --check .` pass
+- [ ] `pyright osmosis_ai/` passes
+- [ ] `pytest` passes (new tests added if applicable)
+- [ ] Public API changes are documented
+- [ ] No secrets or credentials included
diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 00000000..7f43c8a1
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,33 @@
+# Security Policy
+
+## Supported Versions
+
+| Version | Supported |
+|---------|--------------------|
+| Latest | :white_check_mark: |
+| < Latest | :x: |
+
+Only the latest published version of `osmosis-ai` receives security updates.
+
+## Reporting a Vulnerability
+
+**Please do NOT open a public GitHub issue for security vulnerabilities.**
+
+Instead, report vulnerabilities by emailing **brian@osmosis.ai**.
+
+Include the following in your report:
+
+- Description of the vulnerability
+- Steps to reproduce
+- Affected versions
+- Potential impact
+
+## Response Timeline
+
+- **Acknowledgement**: within 3 business days
+- **Initial assessment**: within 7 business days
+- **Fix or mitigation**: depends on severity, targeting 30 days for critical issues
+
+## Disclosure Policy
+
+We follow coordinated disclosure. We ask that you give us reasonable time to address the vulnerability before public disclosure. We will credit reporters in the release notes unless anonymity is requested.
From 6d2e096eb7d12e84a302e50de7502c5a102c0750 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 11:27:01 -0800
Subject: [PATCH 31/60] Enhance PR template and update GitHub Actions workflow
- Added guidelines for PR title formatting and checklist items in the PR template.
- Updated permissions in the publish workflow to allow write access for contents.
- Introduced steps to extract version from `consts.py`, create a Git tag, and generate a GitHub release.
---
.github/PULL_REQUEST_TEMPLATE.md | 13 ++++++++++++
.github/release.yml | 20 ++++++++++++++++++
.github/workflows/check-pr-title.yml | 31 ++++++++++++++++++++++++++++
.github/workflows/publish.yml | 22 +++++++++++++++++++-
4 files changed, 85 insertions(+), 1 deletion(-)
create mode 100644 .github/release.yml
create mode 100644 .github/workflows/check-pr-title.yml
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index ce6a50e2..2ec75e83 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -1,3 +1,14 @@
+
+
## What
@@ -22,6 +33,8 @@
## Checklist
+- [ ] PR title follows `[module] type: description` format
+- [ ] Appropriate labels added (e.g. `enhancement`, `bug`, `breaking`)
- [ ] `ruff check .` and `ruff format --check .` pass
- [ ] `pyright osmosis_ai/` passes
- [ ] `pytest` passes (new tests added if applicable)
diff --git a/.github/release.yml b/.github/release.yml
new file mode 100644
index 00000000..314a5169
--- /dev/null
+++ b/.github/release.yml
@@ -0,0 +1,20 @@
+changelog:
+ categories:
+ - title: "\u26a0\ufe0f Breaking Changes"
+ labels:
+ - breaking
+ - title: "\ud83d\ude80 Features"
+ labels:
+ - enhancement
+ - title: "\ud83d\udc1b Bug Fixes"
+ labels:
+ - bug
+ - title: "\ud83d\udcd6 Documentation"
+ labels:
+ - documentation
+ - title: "\ud83d\udd27 Maintenance"
+ labels:
+ - chore
+ - ci
+ - dependencies
+ - refactor
diff --git a/.github/workflows/check-pr-title.yml b/.github/workflows/check-pr-title.yml
new file mode 100644
index 00000000..d2dbb836
--- /dev/null
+++ b/.github/workflows/check-pr-title.yml
@@ -0,0 +1,31 @@
+name: Check PR Title
+
+on:
+ pull_request:
+ types: [opened, edited, synchronize]
+
+jobs:
+ check-title:
+ runs-on: ubuntu-latest
+ steps:
+ - name: Validate PR title format
+ env:
+ PR_TITLE: ${{ github.event.pull_request.title }}
+ run: |
+ PATTERN='^\[BREAKING\]?\[(reward|rollout|server|cli|auth|eval|misc|ci|doc)\] (feat|fix|refactor|chore|test|doc): .+'
+ if [[ ! "$PR_TITLE" =~ $PATTERN ]]; then
+ echo "::error::PR title does not match required format."
+ echo ""
+ echo "Expected format: [module] type: description"
+ echo " modules: reward, rollout, server, cli, auth, eval, misc, ci, doc"
+ echo " types: feat, fix, refactor, chore, test, doc"
+ echo ""
+ echo "Examples:"
+ echo " [rollout] feat: add streaming support for chat completions"
+ echo " [BREAKING][reward] refactor: rename decorator parameters"
+ echo " [ci] chore: update GitHub Actions versions"
+ echo ""
+ echo "Got: $PR_TITLE"
+ exit 1
+ fi
+ echo "PR title is valid: $PR_TITLE"
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index a9921c8d..0177278c 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -11,7 +11,7 @@ jobs:
environment: Release-secret
runs-on: ubuntu-latest
permissions:
- contents: read
+ contents: write
steps:
- name: Checkout repository
@@ -30,3 +30,23 @@ jobs:
with:
user: __token__
password: ${{ secrets.PYPI_API_TOKEN }}
+
+ - name: Extract version from consts.py
+ id: version
+ run: |
+ VERSION=$(python -c "exec(open('osmosis_ai/consts.py').read()); print(PACKAGE_VERSION)")
+ echo "version=$VERSION" >> "$GITHUB_OUTPUT"
+ echo "tag=v$VERSION" >> "$GITHUB_OUTPUT"
+
+ - name: Create Git tag
+ run: |
+ git tag ${{ steps.version.outputs.tag }}
+ git push origin ${{ steps.version.outputs.tag }}
+
+ - name: Create GitHub Release
+ env:
+ GH_TOKEN: ${{ github.token }}
+ run: |
+ gh release create ${{ steps.version.outputs.tag }} \
+ --title "${{ steps.version.outputs.tag }}" \
+ --generate-notes
From 723d628153cd2864233acb0cacd29570bf04af67 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 11:34:26 -0800
Subject: [PATCH 32/60] update auto-approve
---
.github/workflows/auto-approve.yml | 12 +++++++++++-
1 file changed, 11 insertions(+), 1 deletion(-)
diff --git a/.github/workflows/auto-approve.yml b/.github/workflows/auto-approve.yml
index 31c6328b..c323e985 100644
--- a/.github/workflows/auto-approve.yml
+++ b/.github/workflows/auto-approve.yml
@@ -1,10 +1,20 @@
name: Auto approve
+
on: pull_request_target
jobs:
- build:
+ auto-approve:
runs-on: ubuntu-latest
permissions:
pull-requests: write
+ if: >-
+ github.actor == 'dependabot[bot]' ||
+ github.actor == 'BaiqingL' ||
+ github.actor == 'JoyboyBrian' ||
+ github.actor == 'JakeTrock' ||
+ github.actor == 'allenlinsh' ||
+ github.actor == 'mathewjhan' ||
+ github.actor == 'pandyamarut' ||
+ github.actor == 'artem-osmosis'
steps:
- uses: hmarr/auto-approve-action@v4
From 75e8ed9e91d75bd144b74c136e55a38c909c13f0 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 11:38:35 -0800
Subject: [PATCH 33/60] Enhance contribution guidelines and CI workflow
- Added detailed PR title formatting requirements to CONTRIBUTING.md, including examples and label usage.
- Updated the check-pr-title GitHub Action to enforce the new PR title format with optional breaking change indication.
---
.github/workflows/check-pr-title.yml | 2 +-
CONTRIBUTING.md | 48 ++++++++++++++++++++++++++--
2 files changed, 47 insertions(+), 3 deletions(-)
diff --git a/.github/workflows/check-pr-title.yml b/.github/workflows/check-pr-title.yml
index d2dbb836..893c2aec 100644
--- a/.github/workflows/check-pr-title.yml
+++ b/.github/workflows/check-pr-title.yml
@@ -12,7 +12,7 @@ jobs:
env:
PR_TITLE: ${{ github.event.pull_request.title }}
run: |
- PATTERN='^\[BREAKING\]?\[(reward|rollout|server|cli|auth|eval|misc|ci|doc)\] (feat|fix|refactor|chore|test|doc): .+'
+ PATTERN='^(\[BREAKING\])?\[(reward|rollout|server|cli|auth|eval|misc|ci|doc)\] (feat|fix|refactor|chore|test|doc): .+'
if [[ ! "$PR_TITLE" =~ $PATTERN ]]; then
echo "::error::PR title does not match required format."
echo ""
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 2bb23b8b..0b50d4b9 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -71,9 +71,53 @@ pre-commit install
## Pull Requests
+### PR Title Format
+
+All PR titles **must** follow this format (enforced by CI):
+
+```
+[module] type: description
+```
+
+- **Modules**: `reward`, `rollout`, `server`, `cli`, `auth`, `eval`, `misc`, `ci`, `doc`
+- **Types**: `feat`, `fix`, `refactor`, `chore`, `test`, `doc`
+
+For breaking changes, add `[BREAKING]` before the module:
+
+```
+[BREAKING][module] type: description
+```
+
+**Examples:**
+
+```
+[rollout] feat: add streaming support for chat completions
+[server] fix: handle timeout in rollout init
+[cli] chore: update dependency versions
+[BREAKING][reward] refactor: rename decorator parameters
+```
+
+PR titles appear directly in auto-generated GitHub Release Notes, so keep them clear and descriptive.
+
+### Labels
+
+Add a label to your PR so it gets categorized correctly in Release Notes:
+
+| Label | Use when |
+|-------|----------|
+| `enhancement` | New feature |
+| `bug` | Bug fix |
+| `breaking` | Breaking change |
+| `documentation` | Docs update |
+| `chore` / `ci` / `refactor` / `dependencies` | Maintenance work |
+
+Module-specific labels (`reward`, `rollout`, `server`, `cli`, `auth`, `eval`) can also be added for filtering.
+
+### Workflow
+
1. Fork the repository and create a feature branch
2. Make your changes
3. Run `uv run pytest` and `uv run ruff check .`
-4. Submit a pull request
+4. Submit a pull request with a properly formatted title and label
-CI will run linting, type checking (pyright + mypy), tests across Python 3.10–3.13, and a build validation on every PR.
+CI will run linting, type checking (pyright + mypy), tests across Python 3.10-3.13, PR title validation, and a build validation on every PR.
From 09dcdc8c677f437338f5e96eb50995c4112ea38e Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 12:59:11 -0800
Subject: [PATCH 34/60] Update README and documentation for dataset formats and
CLI options
- Consolidated dataset format information across multiple documents, emphasizing support for Parquet, JSONL, and CSV formats.
- Updated CLI help messages to reflect the recommended dataset file formats.
- Removed redundant sections from the README and examples to streamline content.
- Added new tests for CSV file handling in the DatasetReader class.
---
README.md | 336 ++----------------
docs/README.md | 34 ++
docs/cli.md | 197 ++++++++++
docs/rewards.md | 189 ++++++++++
docs/rollout/README.md | 3 +
docs/rollout/dataset-format.md | 59 +++
docs/rollout/deployment.md | 108 ++++++
docs/rollout/eval.md | 14 +-
docs/rollout/examples.md | 257 +-------------
docs/rollout/test-mode.md | 30 +-
docs/rollout/testing.md | 153 ++++++++
osmosis_ai/rollout/eval/common/dataset.py | 57 +--
osmosis_ai/rollout/eval/evaluation/cli.py | 2 +-
osmosis_ai/rollout/eval/test_mode/cli.py | 2 +-
.../unit/rollout/eval/common/test_dataset.py | 289 +++++++++------
15 files changed, 988 insertions(+), 742 deletions(-)
create mode 100644 docs/README.md
create mode 100644 docs/cli.md
create mode 100644 docs/rewards.md
create mode 100644 docs/rollout/dataset-format.md
create mode 100644 docs/rollout/deployment.md
create mode 100644 docs/rollout/testing.md
diff --git a/README.md b/README.md
index 6a4c933e..aa25da2c 100644
--- a/README.md
+++ b/README.md
@@ -13,11 +13,7 @@ A Python SDK for Osmosis LLM training workflows:
pip install osmosis-ai
```
-Requires Python 3.10 or newer.
-
-This installs the Osmosis CLI and pulls in `litellm` (unified LLM interface supporting 100+ providers) along with supporting utilities such as `PyYAML`, `python-dotenv`, and `requests`.
-
-For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
+Requires Python 3.10 or newer. For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
## Quick Start
@@ -55,320 +51,26 @@ print(rubric_score) # -> 1.0 (full payload available via return_details=True)
## Remote Rollout SDK
-If you're integrating an agent loop with Osmosis remote rollout / TrainGate, see:
-- `docs/rollout/README.md` (quick start)
-- `docs/rollout/architecture.md` (protocol + lifecycle)
-- `docs/rollout/test-mode.md` (local testing with external LLMs)
-- `docs/rollout/eval.md` (model evaluation with eval functions and pass@k metrics)
-
-## Remote Rubric Evaluation
-
-`evaluate_rubric` talks to hosted LLM providers through [LiteLLM](https://github.com/BerriAI/litellm), a unified interface supporting 100+ providers (OpenAI, Anthropic, Google Gemini, xAI, OpenRouter, Cerebras, Azure, Bedrock, Vertex AI, and more). Every provider returns a strict JSON object with `{"score": number, "explanation": string}`. The helper clamps the score into your configured range, validates the structure, and exposes the raw payload when `return_details=True`.
-
-Credentials are resolved from environment variables by default:
-
-- `OPENAI_API_KEY` for OpenAI
-- `ANTHROPIC_API_KEY` for Anthropic
-- `GEMINI_API_KEY` for Google Gemini
-- `XAI_API_KEY` for xAI
-- `OPENROUTER_API_KEY` for OpenRouter
-- `CEREBRAS_API_KEY` for Cerebras
-
-Override the environment variable name with `model_info={"api_key_env": "CUSTOM_ENV_NAME"}` when needed, or supply an inline secret with `model_info={"api_key": "sk-..."}` for ephemeral credentials. Missing API keys raise a `MissingAPIKeyError` that explains how to export the secret before trying again.
-
-`api_key` and `api_key_env` are mutually exclusive ways to provide the same credential. When `api_key` is present and non-empty it is used directly, skipping any environment lookup. Otherwise the resolver falls back to `api_key_env` (or the provider default) and pulls the value from your local environment with `os.getenv`.
-
-`model_info` accepts additional rubric-specific knobs:
-
-- `score_min` / `score_max` – change the default `[0.0, 1.0]` scoring bounds.
-- `system_prompt` / `original_input` – provide optional context strings that will be quoted in the judging prompt.
-- `timeout` – customise the provider timeout in seconds.
-
-Pass `metadata={...}` to `evaluate_rubric` when you need structured context quoted in the judge prompt, and set `return_details=True` to receive the full `RewardRubricRunResult` payload (including the provider’s raw response).
-
-Remote failures surface as `ProviderRequestError` instances, with `ModelNotFoundError` reserved for missing model identifiers so you can retry with a new snapshot.
-
-> Provider model snapshot names change frequently. Check each vendor's dashboard for the latest identifier if you encounter a "model not found" error.
-
-### Provider Architecture
-
-All provider routing is handled by [LiteLLM](https://github.com/BerriAI/litellm). To use any supported provider, pass its name and model in `model_info`:
-
-```python
-result = evaluate_rubric(
- rubric="...",
- solution_str="...",
- model_info={"provider": "anthropic", "model": "claude-sonnet-4-5-20250929"},
-)
-```
-
-Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
-
-## Required Function Signature
-
-All functions decorated with `@osmosis_reward` must have exactly this signature:
-
-```python
-@osmosis_reward
-def your_function(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- # Your reward logic here
- return float_score
-```
-
-### Parameters
-
-- **`solution_str: str`** - The solution string to evaluate (required)
-- **`ground_truth: str`** - The correct/expected answer (required)
-- **`extra_info: dict = None`** - Optional dictionary for additional configuration
+If you're integrating an agent loop with Osmosis remote rollout / TrainGate, see the [Rollout Quick Start](docs/rollout/README.md).
-### Return Value
+## Documentation
-- **`-> float`** - Must return a float value representing the reward score
+| Topic | Description |
+|-------|-------------|
+| [Rewards & Rubrics](docs/rewards.md) | `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric` |
+| [CLI Reference](docs/cli.md) | All `osmosis` commands: auth, serve, test, eval, rubric tools |
+| **Remote Rollout SDK** | |
+| [Quick Start](docs/rollout/README.md) | Get a rollout server running in minutes |
+| [Architecture](docs/rollout/architecture.md) | Protocol design and agent lifecycle |
+| [API Reference](docs/rollout/api-reference.md) | Endpoints, schemas, types |
+| [Examples](docs/rollout/examples.md) | Agent implementations and utilities |
+| [Dataset Format](docs/rollout/dataset-format.md) | Supported formats and required columns |
+| [Test Mode](docs/rollout/test-mode.md) | Test agents with cloud LLMs |
+| [Evaluation](docs/rollout/eval.md) | Evaluate agents with eval functions and pass@k |
+| [Testing](docs/rollout/testing.md) | Unit tests and mock trainer |
+| [Deployment](docs/rollout/deployment.md) | Docker, health checks, production config |
-The decorator will raise a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
-
-## Rubric Function Signature
-
-Rubric functions decorated with `@osmosis_rubric` must match this signature:
-
-```python
-@osmosis_rubric
-def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
- # Your rubric logic here
- return float_score
-```
-
-> The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
-
-### Required `extra_info` fields
-
-- **`provider`** – Non-empty string identifying the judge provider.
-- **`model`** – Non-empty string naming the provider model to call.
-- **`rubric`** – Natural-language rubric instructions for the judge model.
-- **`api_key` / `api_key_env`** – Supply either the raw key or the environment variable name that exposes it.
-
-### Optional `extra_info` fields
-
-- **`system_prompt`** – Optional string prepended to the provider’s base system prompt when invoking the judge; include it inside `extra_info` rather than as a separate argument.
-- **`score_min` / `score_max`** – Optional numeric overrides for the expected score range.
-- **`model_info_overrides`** – Optional dict merged into the provider configuration passed to the judge.
-
-Additional keys are passthrough and can be used for custom configuration. If you need to extend the provider payload (for example adding `api_key_env`), add a dict under `model_info_overrides` and it will be merged with the required `provider`/`model` pair before invoking `evaluate_rubric`. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
-
-> Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
-
-> Tip: When delegating to `evaluate_rubric`, pass the raw `solution_str` directly and include any extra context inside the `metadata` payload.
-
-## Examples
-
-See the [`examples/`](examples/) directory for complete examples:
-
-```python
-@osmosis_reward
-def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- """Case-insensitive string matching with partial credit."""
- match = solution_str.lower().strip() == ground_truth.lower().strip()
-
- if extra_info and 'partial_credit' in extra_info:
- if not match and extra_info['partial_credit']:
- len_diff = abs(len(solution_str) - len(ground_truth))
- if len_diff <= 2:
- return 0.5
-
- return 1.0 if match else 0.0
-
-@osmosis_reward
-def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- """Numeric comparison with configurable tolerance."""
- try:
- solution_num = float(solution_str.strip())
- truth_num = float(ground_truth.strip())
-
- tolerance = extra_info.get('tolerance', 0.01) if extra_info else 0.01
- return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
- except ValueError:
- return 0.0
-```
-
-- `examples/rubric_functions.py` demonstrates `evaluate_rubric` with OpenAI, Anthropic, Gemini, xAI, OpenRouter, and Cerebras via LiteLLM's unified interface.
-- `examples/reward_functions.py` keeps local reward helpers that showcase the decorator contract without external calls.
-- `examples/rubric_configs.yaml` bundles two rubric definitions with provider configuration and scoring bounds.
-- `examples/sample_data.jsonl` contains two rubric-aligned solution strings 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", "original_input": "...", "solution_str": "..."}
-{"conversation_id": "ticket-047", "rubric_id": "policy_grounding", "original_input": "...", "solution_str": "..."}
-```
-
-## CLI Tools
-
-Installing the SDK also provides a lightweight CLI available as `osmosis` (aliases: `osmosis_ai`, `osmosis-ai`).
-
-### Authentication
-
-Log in to Osmosis AI and manage workspace credentials:
-
-```bash
-# Log in to Osmosis AI (opens browser for authentication)
-osmosis login
-
-# Force re-login, clearing existing credentials
-osmosis login --force
-
-# Print the authentication URL without opening browser
-osmosis login --no-browser
-
-# Show current user and all workspaces
-osmosis whoami
-
-# Logout (interactive workspace selection)
-osmosis logout
-
-# Logout from all workspaces
-osmosis logout --all
-
-# Skip confirmation prompt
-osmosis logout -y
-```
-
-Credentials are saved to `~/.config/osmosis/credentials.json` and include workspace information and token expiration.
-
-### Workspace Management
-
-Manage multiple workspaces after logging in:
-
-```bash
-# List all logged-in workspaces
-osmosis workspace list
-
-# Show the current active workspace
-osmosis workspace current
-
-# Switch to a different workspace
-osmosis workspace switch
-```
-
-You can log in to multiple workspaces and switch between them. Each workspace maintains its own credentials and role information.
-
-### Remote Rollout Server
-
-Start a RolloutServer for an agent loop implementation:
-
-```bash
-# Validate agent loop before starting (checks tools, async run, etc.)
-osmosis validate -m my_agent:agent_loop
-
-# Start server with Platform registration (requires `osmosis login`)
-osmosis login
-osmosis serve -m my_agent:agent_loop
-
-# Specify port
-osmosis serve -m my_agent:agent_loop -p 8080
-
-# Local / container mode: skip Platform registration (no login required).
-# NOTE: API key auth is still enabled by default.
-osmosis serve -m my_agent:agent_loop --skip-register
-
-# Local debug mode: disable API key auth AND skip Platform registration
-osmosis serve -m my_agent:agent_loop --local
-
-# Provide a stable API key (otherwise one is generated and printed on startup)
-osmosis serve -m my_agent:agent_loop --skip-register --api-key "$MY_API_KEY"
-
-# Skip validation (not recommended)
-osmosis serve -m my_agent:agent_loop --no-validate
-```
-
-The module path format is `module:attribute`, e.g., `server:agent_loop` or `mypackage.agents:MyAgentClass`.
-
-Note: The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This key is **not** the same as your `osmosis login` token (which is for authenticating with the Osmosis Platform), nor is it used for callbacks *from* your server back to TrainGate.
-
-### Agent Testing
-
-Test your agent loop locally against a dataset using external LLMs:
-
-```bash
-# Batch test with OpenAI
-osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini
-
-# Interactive step-by-step debugging
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
-
-# Git-sync users: test MCP tools directly (no AgentLoop code needed)
-osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-```
-
-See `docs/rollout/test-mode.md` for full documentation.
-
-### Agent Evaluation
-
-Evaluate trained models with custom eval functions and pass@k metrics:
-
-```bash
-# Benchmark a trained model at a serving endpoint
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# pass@k with 5 runs per row
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Git-sync users: evaluate MCP tools directly
-osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model openai/gpt-5-mini
-```
-
-See `docs/rollout/eval.md` for full documentation.
-
-### Rubric Tools
-
-Preview a rubric file and print every configuration discovered, including nested entries:
-
-```bash
-osmosis preview --path path/to/rubric.yaml
-```
-
-Preview a dataset of rubric-scored solutions 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 --rubric support_followup --data examples/sample_data.jsonl
-```
-
-- Command split (development-stage breaking change):
- - `osmosis eval-rubric` evaluates JSONL conversations against hosted rubrics.
- - `osmosis eval` runs rollout eval functions against `RolloutAgentLoop` datasets.
-
-- 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_.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_eval_result_.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.
-- Each dataset record must provide a non-empty `solution_str`; optional fields such as `original_input`, `ground_truth`, and `extra_info` travel with the record and are forwarded to the evaluator when present.
-- When delegating to a custom `@osmosis_rubric` function, the CLI enriches `extra_info` with the active `provider`, `model`, `rubric`, score bounds, any configured `system_prompt`, the resolved `original_input`, and the record’s metadata/extra fields so the decorator’s required entries are always present.
-- Rubric configuration files intentionally reject `extra_info`; provide per-example context through the dataset instead.
-
-Both commands validate the file, echo a short summary (`Loaded ...`), 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.
+Full documentation index: [docs/README.md](docs/README.md)
## Running Examples
@@ -377,6 +79,8 @@ PYTHONPATH=. python examples/reward_functions.py
PYTHONPATH=. python examples/rubric_functions.py # Uncomment the provider you need before running
```
+See the [`examples/`](examples/) directory for all examples.
+
## Contributing
See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, testing, linting, and PR guidelines.
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 00000000..bcb9cbba
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,34 @@
+# Osmosis SDK Documentation
+
+Python SDK for Osmosis AI training workflows, providing reward function validation, rubric evaluation via LLM-as-judge, a Remote Rollout SDK for integrating agent frameworks with Osmosis training, and CLI tools for testing, evaluation, and workspace management.
+
+## Getting Started
+
+- [Installation & Quick Start](../README.md)
+- [Contributing](../CONTRIBUTING.md)
+
+## Rewards & Rubrics
+
+- [Rewards & Rubrics Guide](./rewards.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
+- [Example Code](../examples/) -- reward functions, rubric configs, sample data
+
+## Remote Rollout SDK
+
+- [Quick Start](./rollout/README.md) -- get a rollout server running in minutes
+- [Architecture](./rollout/architecture.md) -- protocol design and agent lifecycle
+- [API Reference](./rollout/api-reference.md) -- endpoints, schemas, types
+- [Examples](./rollout/examples.md) -- agent implementations and utilities
+- [Dataset Format](./rollout/dataset-format.md) -- supported formats and required columns
+- [Test Mode](./rollout/test-mode.md) -- test agents with cloud LLMs
+- [Evaluation](./rollout/eval.md) -- evaluate agents against datasets
+- [Testing](./rollout/testing.md) -- unit tests and mock trainer
+- [Deployment](./rollout/deployment.md) -- Docker, health checks, production config
+
+## CLI Reference
+
+- [CLI Reference](./cli.md) -- all `osmosis` commands
+
+## Other
+
+- [Security Policy](../SECURITY.md)
+- [License](../LICENSE)
diff --git a/docs/cli.md b/docs/cli.md
new file mode 100644
index 00000000..3b78040d
--- /dev/null
+++ b/docs/cli.md
@@ -0,0 +1,197 @@
+# CLI Reference
+
+Installing the SDK provides a lightweight CLI available as `osmosis` (aliases: `osmosis_ai`, `osmosis-ai`). The CLI covers authentication, workspace management, rollout server operation, agent testing and evaluation, and rubric tools. It automatically loads `.env` from the current working directory via `python-dotenv`.
+
+## Authentication
+
+Log in to Osmosis AI and manage credentials. Credentials are saved to `~/.config/osmosis/credentials.json` and include workspace information and token expiration.
+
+### osmosis login
+
+Open a browser-based authentication flow to log in to Osmosis AI:
+
+```bash
+# Log in (opens browser for authentication)
+osmosis login
+
+# Force re-login, clearing existing credentials
+osmosis login --force
+
+# Print the authentication URL without opening browser
+osmosis login --no-browser
+```
+
+### osmosis logout
+
+End your session and revoke stored credentials:
+
+```bash
+# Logout (interactive workspace selection)
+osmosis logout
+
+# Logout from all workspaces
+osmosis logout --all
+
+# Skip confirmation prompt
+osmosis logout -y
+```
+
+### osmosis whoami
+
+Display the current user and all workspaces:
+
+```bash
+osmosis whoami
+```
+
+## Workspace Management
+
+### osmosis workspace
+
+Manage multiple workspaces after logging in. You can log in to multiple workspaces and switch between them. Each workspace maintains its own credentials and role information.
+
+```bash
+# List all logged-in workspaces
+osmosis workspace list
+
+# Show the current active workspace
+osmosis workspace current
+
+# Switch to a different workspace
+osmosis workspace switch
+```
+
+## Remote Rollout Server
+
+### osmosis serve
+
+Start a RolloutServer for an agent loop implementation. The module path format is `module:attribute`, e.g., `server:agent_loop` or `mypackage.agents:MyAgentClass`.
+
+```bash
+# Start server with Platform registration (requires `osmosis login`)
+osmosis login
+osmosis serve -m my_agent:agent_loop
+
+# Specify port
+osmosis serve -m my_agent:agent_loop -p 8080
+
+# Local / container mode: skip Platform registration (no login required).
+# NOTE: API key auth is still enabled by default.
+osmosis serve -m my_agent:agent_loop --skip-register
+
+# Local debug mode: disable API key auth AND skip Platform registration
+osmosis serve -m my_agent:agent_loop --local
+
+# Provide a stable API key (otherwise one is generated and printed on startup)
+osmosis serve -m my_agent:agent_loop --skip-register --api-key "$MY_API_KEY"
+
+# Skip validation (not recommended)
+osmosis serve -m my_agent:agent_loop --no-validate
+```
+
+Note: The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This key is **not** the same as your `osmosis login` token (which is for authenticating with the Osmosis Platform), nor is it used for callbacks *from* your server back to TrainGate.
+
+### osmosis validate
+
+Validate an agent loop before starting the server (checks tools, async run method, etc.):
+
+```bash
+osmosis validate -m my_agent:agent_loop
+```
+
+## Agent Testing
+
+### osmosis test
+
+Test your agent loop locally against a dataset using external LLMs:
+
+```bash
+# Batch test with OpenAI
+osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini
+
+# Interactive step-by-step debugging
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
+
+# Git-sync users: test MCP tools directly (no AgentLoop code needed)
+osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
+```
+
+See [Test Mode](./rollout/test-mode.md) for full documentation on dataset format, interactive mode, and MCP tool testing.
+
+## Agent Evaluation
+
+### osmosis eval
+
+Evaluate trained models with custom eval functions and pass@k metrics:
+
+```bash
+# Benchmark a trained model at a serving endpoint
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
+ --eval-fn rewards:compute_reward \
+ --model my-finetuned-model --base-url http://localhost:8000/v1
+
+# pass@k with 5 runs per row
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
+ --eval-fn rewards:compute_reward --n 5 \
+ --model my-finetuned-model --base-url http://localhost:8000/v1
+
+# Git-sync users: evaluate MCP tools directly
+osmosis eval --mcp ./mcp -d data.jsonl \
+ --eval-fn rewards:compute_reward \
+ --model openai/gpt-5-mini
+```
+
+See [Evaluation](./rollout/eval.md) for full documentation on eval functions, pass@k metrics, and output formats.
+
+## Rubric Tools
+
+### osmosis preview
+
+Preview a rubric file and print every configuration discovered, including nested entries:
+
+```bash
+osmosis preview --path path/to/rubric.yaml
+```
+
+Preview a dataset of rubric-scored solutions stored as JSONL:
+
+```bash
+osmosis preview --path path/to/data.jsonl
+```
+
+Both formats validate the file, echo a short summary (`Loaded ...`), 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.
+
+### osmosis eval-rubric
+
+Evaluate a dataset against a hosted rubric configuration and print the returned scores:
+
+```bash
+osmosis eval-rubric --rubric support_followup --data examples/sample_data.jsonl
+```
+
+**Command split** (development-stage breaking change):
+- `osmosis eval-rubric` evaluates JSONL conversations against hosted rubrics.
+- `osmosis eval` runs rollout eval functions against `RolloutAgentLoop` datasets.
+
+**Options:**
+
+- `-d`/`--data path/to/data.jsonl` -- Supply the dataset; the path is resolved relative to the current working directory.
+- `--config path/to/rubric_configs.yaml` -- Provide rubric definitions when they are not located alongside the dataset.
+- `-n`/`--number` -- Sample the provider multiple times per record; the CLI prints every run along with aggregate statistics (average, variance, standard deviation, and min/max).
+- `--output path/to/dir` -- Create the directory (if needed) and emit `rubric_eval_result_.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.
+
+**Behavior notes:**
+
+- Skip `--output` to collect results under `~/.cache/osmosis/eval_result//rubric_eval_result_.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.
+- Each dataset record must provide a non-empty `solution_str`; optional fields such as `original_input`, `ground_truth`, and `extra_info` travel with the record and are forwarded to the evaluator when present.
+- When delegating to a custom `@osmosis_rubric` function, the CLI enriches `extra_info` with the active `provider`, `model`, `rubric`, score bounds, any configured `system_prompt`, the resolved `original_input`, and the record's metadata/extra fields so the decorator's required entries are always present.
+- Rubric configuration files intentionally reject `extra_info`; provide per-example context through the dataset instead.
+
+See [Rewards & Rubrics](./rewards.md) for details on `@osmosis_reward`, `@osmosis_rubric`, and `evaluate_rubric`.
+
+## See Also
+
+- [Rewards & Rubrics](./rewards.md)
+- [Rollout Test Mode](./rollout/test-mode.md)
+- [Rollout Evaluation](./rollout/eval.md)
diff --git a/docs/rewards.md b/docs/rewards.md
new file mode 100644
index 00000000..624ff1ab
--- /dev/null
+++ b/docs/rewards.md
@@ -0,0 +1,189 @@
+# Rewards & Rubrics
+
+The Osmosis SDK provides two mechanisms for scoring LLM outputs during training: **reward functions** for deterministic, code-based scoring, and **rubric evaluations** for LLM-as-judge assessments powered by external providers. Both integrate with the Osmosis training platform to drive reinforcement learning workflows.
+
+## @osmosis_reward
+
+All functions decorated with `@osmosis_reward` must have exactly this signature:
+
+```python
+from osmosis_ai import osmosis_reward
+
+@osmosis_reward
+def your_function(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+ # Your reward logic here
+ return float_score
+```
+
+### Parameters
+
+- **`solution_str: str`** -- The solution string to evaluate (required).
+- **`ground_truth: str`** -- The correct/expected answer (required).
+- **`extra_info: dict = None`** -- Optional dictionary for additional configuration.
+
+### Return Value
+
+- **`-> float`** -- Must return a float value representing the reward score.
+
+The decorator will raise a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
+
+## @osmosis_rubric
+
+Rubric functions decorated with `@osmosis_rubric` must match this signature:
+
+```python
+from osmosis_ai import osmosis_rubric
+
+@osmosis_rubric
+def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+ # Your rubric logic here
+ return float_score
+```
+
+> The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
+
+### Required `extra_info` fields
+
+- **`provider`** -- Non-empty string identifying the judge provider.
+- **`model`** -- Non-empty string naming the provider model to call.
+- **`rubric`** -- Natural-language rubric instructions for the judge model.
+- **`api_key` / `api_key_env`** -- Supply either the raw key or the environment variable name that exposes it.
+
+### Optional `extra_info` fields
+
+- **`system_prompt`** -- Optional string prepended to the provider's base system prompt when invoking the judge; include it inside `extra_info` rather than as a separate argument.
+- **`score_min` / `score_max`** -- Optional numeric overrides for the expected score range.
+- **`model_info_overrides`** -- Optional dict merged into the provider configuration passed to the judge.
+
+Additional keys are passthrough and can be used for custom configuration. If you need to extend the provider payload (for example adding `api_key_env`), add a dict under `model_info_overrides` and it will be merged with the required `provider`/`model` pair before invoking `evaluate_rubric`. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
+
+> Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
+
+> Tip: When delegating to `evaluate_rubric`, pass the raw `solution_str` directly and include any extra context inside the `metadata` payload.
+
+## evaluate_rubric
+
+`evaluate_rubric` talks to hosted LLM providers through [LiteLLM](https://github.com/BerriAI/litellm), a unified interface supporting 100+ providers (OpenAI, Anthropic, Google Gemini, xAI, OpenRouter, Cerebras, Azure, Bedrock, Vertex AI, and more). Every provider returns a strict JSON object with `{"score": number, "explanation": string}`. The helper clamps the score into your configured range, validates the structure, and exposes the raw payload when `return_details=True`.
+
+### Basic Usage
+
+```python
+from osmosis_ai import evaluate_rubric
+
+solution = "The capital of France is Paris."
+
+rubric_score = evaluate_rubric(
+ rubric="Assistant must mention the verified capital city.",
+ solution_str=solution,
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5",
+ "api_key_env": "OPENAI_API_KEY",
+ },
+ ground_truth="Paris",
+)
+
+print(rubric_score) # -> 1.0 (full payload available via return_details=True)
+```
+
+### Credentials
+
+Credentials are resolved from environment variables by default:
+
+- `OPENAI_API_KEY` for OpenAI
+- `ANTHROPIC_API_KEY` for Anthropic
+- `GEMINI_API_KEY` for Google Gemini
+- `XAI_API_KEY` for xAI
+- `OPENROUTER_API_KEY` for OpenRouter
+- `CEREBRAS_API_KEY` for Cerebras
+
+Override the environment variable name with `model_info={"api_key_env": "CUSTOM_ENV_NAME"}` when needed, or supply an inline secret with `model_info={"api_key": "sk-..."}` for ephemeral credentials. Missing API keys raise a `MissingAPIKeyError` that explains how to export the secret before trying again.
+
+`api_key` and `api_key_env` are mutually exclusive ways to provide the same credential. When `api_key` is present and non-empty it is used directly, skipping any environment lookup. Otherwise the resolver falls back to `api_key_env` (or the provider default) and pulls the value from your local environment with `os.getenv`.
+
+### Configuration Options
+
+`model_info` accepts additional rubric-specific knobs:
+
+- **`score_min` / `score_max`** -- Change the default `[0.0, 1.0]` scoring bounds.
+- **`system_prompt` / `original_input`** -- Provide optional context strings that will be quoted in the judging prompt.
+- **`timeout`** -- Customise the provider timeout in seconds.
+
+Pass `metadata={...}` to `evaluate_rubric` when you need structured context quoted in the judge prompt, and set `return_details=True` to receive the full `RewardRubricRunResult` payload (including the provider's raw response).
+
+### Provider Architecture
+
+All provider routing is handled by [LiteLLM](https://github.com/BerriAI/litellm). To use any supported provider, pass its name and model in `model_info`:
+
+```python
+result = evaluate_rubric(
+ rubric="...",
+ solution_str="...",
+ model_info={"provider": "anthropic", "model": "claude-sonnet-4-5-20250929"},
+)
+```
+
+Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
+
+### Error Handling
+
+Remote failures surface as `ProviderRequestError` instances, with `ModelNotFoundError` reserved for missing model identifiers so you can retry with a new snapshot.
+
+> Provider model snapshot names change frequently. Check each vendor's dashboard for the latest identifier if you encounter a "model not found" error.
+
+## Examples
+
+See the [`examples/`](../examples/) directory for complete examples:
+
+```python
+@osmosis_reward
+def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+ """Case-insensitive string matching with partial credit."""
+ match = solution_str.lower().strip() == ground_truth.lower().strip()
+
+ if extra_info and 'partial_credit' in extra_info:
+ if not match and extra_info['partial_credit']:
+ len_diff = abs(len(solution_str) - len(ground_truth))
+ if len_diff <= 2:
+ return 0.5
+
+ return 1.0 if match else 0.0
+
+@osmosis_reward
+def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+ """Numeric comparison with configurable tolerance."""
+ try:
+ solution_num = float(solution_str.strip())
+ truth_num = float(ground_truth.strip())
+
+ tolerance = extra_info.get('tolerance', 0.01) if extra_info else 0.01
+ return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
+ except ValueError:
+ return 0.0
+```
+
+- `examples/reward_functions.py` keeps local reward helpers that showcase the decorator contract without external calls.
+- `examples/rubric_functions.py` demonstrates `evaluate_rubric` with OpenAI, Anthropic, Gemini, xAI, OpenRouter, and Cerebras via LiteLLM's unified interface.
+- `examples/rubric_configs.yaml` bundles two rubric definitions with provider configuration and scoring bounds.
+- `examples/sample_data.jsonl` contains two rubric-aligned solution strings 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", "original_input": "...", "solution_str": "..."}
+{"conversation_id": "ticket-047", "rubric_id": "policy_grounding", "original_input": "...", "solution_str": "..."}
+```
+
+## See Also
+
+- [CLI Reference](./cli.md) -- `osmosis preview` and `osmosis eval-rubric` commands
+- [Rollout Evaluation](./rollout/eval.md) -- evaluating agents against datasets
diff --git a/docs/rollout/README.md b/docs/rollout/README.md
index ea78a363..a78d6878 100644
--- a/docs/rollout/README.md
+++ b/docs/rollout/README.md
@@ -228,3 +228,6 @@ export OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS=200
- [API Reference](./api-reference.md) - Complete API documentation
- [Examples](./examples.md) - Full working examples
- [Test Mode](./test-mode.md) - Local testing with external LLM providers
+- [Testing](./testing.md) - Unit tests and mock trainer
+- [Deployment](./deployment.md) - Docker, health checks, production config
+- [Dataset Format](./dataset-format.md) - Supported formats and required columns
diff --git a/docs/rollout/dataset-format.md b/docs/rollout/dataset-format.md
new file mode 100644
index 00000000..d6f3ade0
--- /dev/null
+++ b/docs/rollout/dataset-format.md
@@ -0,0 +1,59 @@
+# Dataset Format
+
+Datasets define the test cases for evaluating and testing your rollout agents. They provide the prompts, expected outputs, and optional metadata used by both [test mode](./test-mode.md) and [evaluation](./eval.md).
+
+## Supported Formats
+
+- **Parquet** (recommended) -- columnar format, compact and fast for large datasets
+- **JSONL** -- one JSON object per line
+- **CSV** -- comma-separated values with a header row
+
+## Required Columns
+
+Each row must contain these columns (case-insensitive):
+
+| Column | Type | Description |
+|--------|------|-------------|
+| `ground_truth` | `str` | Expected output (for reward computation) |
+| `user_prompt` | `str` | User message to start the conversation |
+| `system_prompt` | `str` | System prompt for the LLM |
+
+## Example Dataset (Parquet)
+
+```python
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+table = pa.table({
+ "system_prompt": ["You are a helpful calculator.", "You are a helpful calculator."],
+ "user_prompt": ["What is 2 + 2?", "What is 10 * 5?"],
+ "ground_truth": ["4", "50"],
+})
+pq.write_table(table, "data.parquet")
+```
+
+## Example Dataset (JSONL)
+
+```jsonl
+{"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 2 + 2?", "ground_truth": "4"}
+{"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 10 * 5?", "ground_truth": "50"}
+```
+
+## Example Dataset (CSV)
+
+```csv
+system_prompt,user_prompt,ground_truth
+You are a helpful calculator.,What is 2 + 2?,4
+You are a helpful calculator.,What is 10 * 5?,50
+```
+
+> **Note:** Fields containing commas, newlines, or double-quotes must be enclosed in double-quotes per [RFC 4180](https://tools.ietf.org/html/rfc4180). For datasets where prompts contain rich text, Parquet or JSONL are better choices.
+
+## Additional Columns
+
+Any columns beyond the three required ones are passed to `RolloutRequest.metadata`. This lets you include extra context -- such as difficulty level, category tags, or reference data -- that your agent or reward function can access at runtime.
+
+## See Also
+
+- [Test Mode](./test-mode.md) -- test agents locally with external LLMs
+- [Evaluation](./eval.md) -- evaluate agents with eval functions and pass@k
diff --git a/docs/rollout/deployment.md b/docs/rollout/deployment.md
new file mode 100644
index 00000000..2eb3d528
--- /dev/null
+++ b/docs/rollout/deployment.md
@@ -0,0 +1,108 @@
+# Production Deployment
+
+This guide covers deploying your rollout server to production environments using Docker, Docker Compose, health checks, and logging best practices.
+
+## Docker Configuration
+
+Using the CLI (recommended):
+
+```dockerfile
+# Dockerfile
+
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 9000
+
+# Using CLI (validates on startup)
+CMD ["osmosis", "serve", "-m", "main:agent_loop", "-p", "9000", "--skip-register"]
+```
+
+Or with uvicorn directly:
+
+```dockerfile
+# Dockerfile (alternative)
+
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 9000
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "9000"]
+```
+
+```txt
+# requirements.txt
+osmosis-ai[server,config]>=0.1.0
+```
+
+## Docker Compose
+
+```yaml
+# docker-compose.yml
+
+version: '3.8'
+
+services:
+ rollout-server:
+ build: .
+ ports:
+ - "9000:9000"
+ environment:
+ - OSMOSIS_ROLLOUT_LOG_LEVEL=INFO
+ healthcheck:
+ test: ["CMD", "curl", "-f", "http://localhost:9000/health"]
+ interval: 30s
+ timeout: 10s
+ retries: 3
+```
+
+## Health Check Integration
+
+```python
+# For Kubernetes or load balancer health checks
+
+@app.get("/health")
+async def health():
+ return {
+ "status": "healthy",
+ "agent_loop": agent_loop.name,
+ }
+
+@app.get("/ready")
+async def ready():
+ # Add any readiness checks here
+ return {"ready": True}
+```
+
+## Logging Configuration
+
+```python
+import logging
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+
+# Reduce httpx noise
+logging.getLogger("httpx").setLevel(logging.WARNING)
+```
+
+## See Also
+
+- [Examples](./examples.md) -- agent implementations
+- [Testing](./testing.md) -- unit tests and mock trainer
+- [Architecture](./architecture.md) -- protocol design
diff --git a/docs/rollout/eval.md b/docs/rollout/eval.md
index dcf08013..2cb7de30 100644
--- a/docs/rollout/eval.md
+++ b/docs/rollout/eval.md
@@ -17,7 +17,7 @@ Key capabilities:
- **Concurrent execution** with `--batch-size` for faster benchmarks
- Use LiteLLM providers (e.g., `openai/gpt-5-mini`) as the primary model
-> Command split (development-stage breaking change):
+> Command split (development-stage breaking change):
> `osmosis eval-rubric` is for hosted rubric evaluation, while `osmosis eval` is for agent eval mode documented here.
## Quick Start
@@ -226,15 +226,7 @@ The `--model` parameter should match the model name as registered in the serving
## Dataset Format
-Eval mode uses the same dataset format as [Test Mode](./test-mode.md#dataset-format). Each row must contain these columns (case-insensitive):
-
-| Column | Type | Description |
-|--------|------|-------------|
-| `ground_truth` | `str` | Expected output (passed to eval functions) |
-| `user_prompt` | `str` | User message to start the conversation |
-| `system_prompt` | `str` | System prompt for the LLM |
-
-Additional columns are passed to eval functions via `metadata` (full mode) or `extra_info` (simple mode).
+See [Dataset Format](./dataset-format.md) for supported formats and required columns.
---
@@ -319,7 +311,7 @@ osmosis eval [OPTIONS]
| Option | Description |
|--------|-------------|
-| `-d, --dataset FILE` | Path to dataset file (.json, .jsonl, .parquet) |
+| `-d, --dataset FILE` | Path to dataset file (.parquet recommended, .jsonl, .csv) |
| `--eval-fn MODULE:FN` | Eval function path (can be specified multiple times) |
| `--model MODEL` | Model to benchmark (see [Model Options](#model-options)) |
diff --git a/docs/rollout/examples.md b/docs/rollout/examples.md
index 3002a139..d13444a5 100644
--- a/docs/rollout/examples.md
+++ b/docs/rollout/examples.md
@@ -659,257 +659,6 @@ counts = count_messages_by_role(messages)
---
-## Testing Your Agent
-
-### Using Mock Trainer
-
-The SDK provides a mock trainer for local testing without a real TrainGate server.
-
-```python
-# test_with_mock_trainer.py
-
-import pytest
-from fastapi.testclient import TestClient
-
-from osmosis_ai.rollout.testing import (
- create_mock_trainer_app,
- RolloutCompletionTracker,
- patch_httpx_for_mock_trainer,
-)
-
-
-@pytest.fixture
-def mock_trainer(monkeypatch):
- """Set up mock trainer with completion tracking."""
- tracker = RolloutCompletionTracker()
- app = create_mock_trainer_app(tracker=tracker)
- client = TestClient(app)
- patch_httpx_for_mock_trainer(client, monkeypatch)
- return client, tracker
-
-
-def test_rollout_with_mock_trainer(mock_trainer):
- """Test complete rollout flow with mock trainer."""
- client, tracker = mock_trainer
-
- # Your rollout will use the mock trainer
- # ...
-
- # Wait for completion callback
- assert tracker.wait(timeout=5.0)
- assert len(tracker.responses) == 1
- assert tracker.responses[0]["status"] == "COMPLETED"
-```
-
-### Custom Tool Call Generator
-
-You can customize when the mock trainer generates tool calls:
-
-```python
-def weather_tool_generator(message):
- """Generate weather tool calls for weather-related messages."""
- if "weather" in message.get("content", "").lower():
- return [
- {
- "id": "call_weather",
- "type": "function",
- "function": {"name": "get_weather", "arguments": "{}"},
- }
- ]
- return None
-
-app = create_mock_trainer_app(tool_call_generator=weather_tool_generator)
-```
-
-### Example Test Using FastAPI TestClient
-
-```python
-# test_my_agent.py
-
-import pytest
-from unittest.mock import AsyncMock, patch
-
-from fastapi.testclient import TestClient
-from osmosis_ai.rollout import create_app
-
-from my_agent import MyAgentLoop
-
-
-@pytest.fixture
-def app():
- """Create test application."""
- return create_app(MyAgentLoop())
-
-
-@pytest.fixture
-def client(app):
- """Create test client."""
- with TestClient(app) as client:
- yield client
-
-
-def test_health_endpoint(client):
- """Test health check."""
- response = client.get("/health")
- assert response.status_code == 200
- assert response.json()["status"] == "healthy"
-
-
-def test_init_returns_tools(client):
- """Test /v1/rollout/init returns tools."""
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "test-123",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
- assert response.status_code == 202
- data = response.json()
- assert "tools" in data
-
-
-@pytest.mark.asyncio
-async def test_agent_run():
- """Test agent run directly."""
- from unittest.mock import MagicMock
- from osmosis_ai.rollout import RolloutContext, RolloutRequest, RolloutMetrics
-
- agent = MyAgentLoop()
-
- # Create mock context
- mock_llm = MagicMock()
- mock_llm.get_metrics.return_value = RolloutMetrics()
- mock_llm.chat_completions = AsyncMock(return_value=MagicMock(
- message={"role": "assistant", "content": "Hello!"},
- has_tool_calls=False,
- ))
-
- request = RolloutRequest(
- rollout_id="test",
- server_url="http://localhost",
- messages=[{"role": "user", "content": "Hi"}],
- completion_params={},
- )
-
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- )
-
- result = await agent.run(ctx)
-
- assert result.status == "COMPLETED"
-```
-
----
-
-## Production Deployment
-
-### Docker Configuration
-
-Using the CLI (recommended):
-
-```dockerfile
-# Dockerfile
-
-FROM python:3.11-slim
-
-WORKDIR /app
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-COPY . .
-
-EXPOSE 9000
-
-# Using CLI (validates on startup)
-CMD ["osmosis", "serve", "-m", "main:agent_loop", "-p", "9000", "--skip-register"]
-```
-
-Or with uvicorn directly:
-
-```dockerfile
-# Dockerfile (alternative)
-
-FROM python:3.11-slim
-
-WORKDIR /app
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-COPY . .
-
-EXPOSE 9000
-
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "9000"]
-```
-
-```txt
-# requirements.txt
-osmosis-ai[server,config]>=0.1.0
-```
-
-### Docker Compose
-
-```yaml
-# docker-compose.yml
-
-version: '3.8'
-
-services:
- rollout-server:
- build: .
- ports:
- - "9000:9000"
- environment:
- - OSMOSIS_ROLLOUT_LOG_LEVEL=INFO
- healthcheck:
- test: ["CMD", "curl", "-f", "http://localhost:9000/health"]
- interval: 30s
- timeout: 10s
- retries: 3
-```
-
-### Health Check Integration
-
-```python
-# For Kubernetes or load balancer health checks
-
-@app.get("/health")
-async def health():
- return {
- "status": "healthy",
- "agent_loop": agent_loop.name,
- }
-
-@app.get("/ready")
-async def ready():
- # Add any readiness checks here
- return {"ready": True}
-```
-
-### Logging Configuration
-
-```python
-import logging
-
-logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-
-# Reduce httpx noise
-logging.getLogger("httpx").setLevel(logging.WARNING)
-```
-
----
-
## Using the Registry
For applications with multiple agent types.
@@ -934,3 +683,9 @@ print(list_agent_loops()) # ['calculator', 'search', 'weather']
agent = get_agent_loop("calculator")
app = create_app(agent)
```
+
+## See Also
+
+- [Testing](./testing.md) — unit tests and mock trainer
+- [Deployment](./deployment.md) — Docker, health checks, production config
+- [API Reference](./api-reference.md) — endpoints, schemas, types
diff --git a/docs/rollout/test-mode.md b/docs/rollout/test-mode.md
index e959fa30..31fb3009 100644
--- a/docs/rollout/test-mode.md
+++ b/docs/rollout/test-mode.md
@@ -105,31 +105,7 @@ print(f"Total tokens: {results.total_tokens}")
## Dataset Format
-Test mode reads datasets in JSON, JSONL, or Parquet format. Each row must contain these columns (case-insensitive):
-
-| Column | Type | Description |
-|--------|------|-------------|
-| `ground_truth` | `str` | Expected output (for reward computation) |
-| `user_prompt` | `str` | User message to start the conversation |
-| `system_prompt` | `str` | System prompt for the LLM |
-
-Additional columns are passed to `RolloutRequest.metadata`.
-
-### Example Dataset (JSONL)
-
-```jsonl
-{"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 2 + 2?", "ground_truth": "4"}
-{"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 10 * 5?", "ground_truth": "50"}
-```
-
-### Example Dataset (JSON)
-
-```json
-[
- {"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 2 + 2?", "ground_truth": "4"},
- {"system_prompt": "You are a helpful calculator.", "user_prompt": "What is 10 * 5?", "ground_truth": "50"}
-]
-```
+See [Dataset Format](./dataset-format.md) for supported formats and required columns.
---
@@ -143,7 +119,7 @@ osmosis test [OPTIONS]
| Option | Description |
|--------|-------------|
-| `-d, --dataset FILE` | Path to dataset file (.json, .jsonl, .parquet) |
+| `-d, --dataset FILE` | Path to dataset file (.parquet recommended, .jsonl, .csv) |
### Agent Options (one required, mutually exclusive)
@@ -616,4 +592,4 @@ See [LiteLLM Environment Variables](https://docs.litellm.ai/docs/providers) for
- [Architecture](./architecture.md) - System design overview
- [API Reference](./api-reference.md) - Complete SDK API documentation
- [Examples](./examples.md) - Working code examples
-- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
\ No newline at end of file
+- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
diff --git a/docs/rollout/testing.md b/docs/rollout/testing.md
new file mode 100644
index 00000000..d580803d
--- /dev/null
+++ b/docs/rollout/testing.md
@@ -0,0 +1,153 @@
+# Testing Your Agent
+
+This guide covers unit testing and local testing strategies for rollout agents. These techniques let you validate your agent logic without deploying to training infrastructure or connecting to external LLM providers.
+
+## Using Mock Trainer
+
+The SDK provides a mock trainer for local testing without a real TrainGate server.
+
+```python
+# test_with_mock_trainer.py
+
+import pytest
+from fastapi.testclient import TestClient
+
+from osmosis_ai.rollout.testing import (
+ create_mock_trainer_app,
+ RolloutCompletionTracker,
+ patch_httpx_for_mock_trainer,
+)
+
+
+@pytest.fixture
+def mock_trainer(monkeypatch):
+ """Set up mock trainer with completion tracking."""
+ tracker = RolloutCompletionTracker()
+ app = create_mock_trainer_app(tracker=tracker)
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+ return client, tracker
+
+
+def test_rollout_with_mock_trainer(mock_trainer):
+ """Test complete rollout flow with mock trainer."""
+ client, tracker = mock_trainer
+
+ # Your rollout will use the mock trainer
+ # ...
+
+ # Wait for completion callback
+ assert tracker.wait(timeout=5.0)
+ assert len(tracker.responses) == 1
+ assert tracker.responses[0]["status"] == "COMPLETED"
+```
+
+## Custom Tool Call Generator
+
+You can customize when the mock trainer generates tool calls:
+
+```python
+def weather_tool_generator(message):
+ """Generate weather tool calls for weather-related messages."""
+ if "weather" in message.get("content", "").lower():
+ return [
+ {
+ "id": "call_weather",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": "{}"},
+ }
+ ]
+ return None
+
+app = create_mock_trainer_app(tool_call_generator=weather_tool_generator)
+```
+
+## Example Test Using FastAPI TestClient
+
+```python
+# test_my_agent.py
+
+import pytest
+from unittest.mock import AsyncMock, patch
+
+from fastapi.testclient import TestClient
+from osmosis_ai.rollout import create_app
+
+from my_agent import MyAgentLoop
+
+
+@pytest.fixture
+def app():
+ """Create test application."""
+ return create_app(MyAgentLoop())
+
+
+@pytest.fixture
+def client(app):
+ """Create test client."""
+ with TestClient(app) as client:
+ yield client
+
+
+def test_health_endpoint(client):
+ """Test health check."""
+ response = client.get("/health")
+ assert response.status_code == 200
+ assert response.json()["status"] == "healthy"
+
+
+def test_init_returns_tools(client):
+ """Test /v1/rollout/init returns tools."""
+ response = client.post(
+ "/v1/rollout/init",
+ json={
+ "rollout_id": "test-123",
+ "server_url": "http://localhost:8080",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "completion_params": {"temperature": 0.7},
+ },
+ )
+ assert response.status_code == 202
+ data = response.json()
+ assert "tools" in data
+
+
+@pytest.mark.asyncio
+async def test_agent_run():
+ """Test agent run directly."""
+ from unittest.mock import MagicMock
+ from osmosis_ai.rollout import RolloutContext, RolloutRequest, RolloutMetrics
+
+ agent = MyAgentLoop()
+
+ # Create mock context
+ mock_llm = MagicMock()
+ mock_llm.get_metrics.return_value = RolloutMetrics()
+ mock_llm.chat_completions = AsyncMock(return_value=MagicMock(
+ message={"role": "assistant", "content": "Hello!"},
+ has_tool_calls=False,
+ ))
+
+ request = RolloutRequest(
+ rollout_id="test",
+ server_url="http://localhost",
+ messages=[{"role": "user", "content": "Hi"}],
+ completion_params={},
+ )
+
+ ctx = RolloutContext(
+ request=request,
+ tools=[],
+ llm=mock_llm,
+ )
+
+ result = await agent.run(ctx)
+
+ assert result.status == "COMPLETED"
+```
+
+## See Also
+
+- [Examples](./examples.md) -- agent implementations and utilities
+- [Test Mode](./test-mode.md) -- testing with cloud LLMs
+- [Deployment](./deployment.md) -- production deployment
diff --git a/osmosis_ai/rollout/eval/common/dataset.py b/osmosis_ai/rollout/eval/common/dataset.py
index 228fa06c..6d2bc39f 100644
--- a/osmosis_ai/rollout/eval/common/dataset.py
+++ b/osmosis_ai/rollout/eval/common/dataset.py
@@ -7,6 +7,7 @@
from __future__ import annotations
+import csv
import json
import logging
from collections.abc import Iterator
@@ -35,7 +36,7 @@ class DatasetRow(TypedDict):
class DatasetReader:
- """Reader for JSON / JSONL / Parquet datasets used by local workflows."""
+ """Reader for Parquet / JSONL / CSV datasets used by local workflows."""
def __init__(self, file_path: str) -> None:
self.file_path = Path(file_path)
@@ -44,10 +45,10 @@ def __init__(self, file_path: str) -> None:
raise FileNotFoundError(f"Dataset file not found: {file_path}")
self._extension = self.file_path.suffix.lower()
- if self._extension not in (".json", ".jsonl", ".parquet"):
+ if self._extension not in (".parquet", ".jsonl", ".csv"):
raise DatasetParseError(
f"Unsupported file format: {self._extension}. "
- f"Supported formats: .json, .jsonl, .parquet"
+ f"Supported formats: .parquet (recommended), .jsonl, .csv"
)
self._row_count: int | None = None
@@ -83,44 +84,28 @@ def iter_rows(
count += 1
def _iter_raw_rows(self) -> Iterator[dict[str, Any]]:
- if self._extension == ".json":
- yield from self._parse_json()
- elif self._extension == ".jsonl":
+ if self._extension == ".jsonl":
yield from self._iter_jsonl()
elif self._extension == ".parquet":
yield from self._parse_parquet()
+ elif self._extension == ".csv":
+ yield from self._iter_csv()
def __len__(self) -> int:
if self._row_count is not None:
return self._row_count
- if self._extension == ".json":
- self._row_count = len(self._parse_json())
- elif self._extension == ".jsonl":
+ if self._extension == ".jsonl":
self._row_count = self._count_jsonl_rows()
elif self._extension == ".parquet":
self._row_count = self._count_parquet_rows()
+ elif self._extension == ".csv":
+ self._row_count = self._count_csv_rows()
else:
self._row_count = 0
return self._row_count
- def _parse_json(self) -> list[dict[str, Any]]:
- try:
- with open(self.file_path, encoding="utf-8") as f:
- data = json.load(f)
- except json.JSONDecodeError as e:
- raise DatasetParseError(f"Invalid JSON file: {e}") from e
- except OSError as e:
- raise DatasetParseError(f"Error reading file: {e}") from e
-
- if not isinstance(data, list):
- raise DatasetParseError(
- f"JSON file must contain an array of objects, got {type(data).__name__}"
- )
-
- return data
-
def _iter_jsonl(self) -> Iterator[dict[str, Any]]:
try:
with open(self.file_path, encoding="utf-8") as f:
@@ -179,6 +164,28 @@ def _count_parquet_rows(self) -> int:
except Exception as e:
raise DatasetParseError(f"Error reading Parquet metadata: {e}") from e
+ def _iter_csv(self) -> Iterator[dict[str, Any]]:
+ try:
+ with open(self.file_path, encoding="utf-8", newline="") as f:
+ reader = csv.DictReader(f)
+ for row in reader:
+ yield dict(row)
+ except csv.Error as e:
+ raise DatasetParseError(f"Invalid CSV file: {e}") from e
+ except OSError as e:
+ raise DatasetParseError(f"Error reading file: {e}") from e
+
+ def _count_csv_rows(self) -> int:
+ count = 0
+ try:
+ with open(self.file_path, encoding="utf-8", newline="") as f:
+ reader = csv.DictReader(f)
+ for _ in reader:
+ count += 1
+ except (OSError, csv.Error):
+ pass
+ return count
+
def _validate_row(self, row: Any, row_index: int) -> DatasetRow:
if not isinstance(row, dict):
raise DatasetValidationError(
diff --git a/osmosis_ai/rollout/eval/evaluation/cli.py b/osmosis_ai/rollout/eval/evaluation/cli.py
index 1f1078e9..30a485bb 100644
--- a/osmosis_ai/rollout/eval/evaluation/cli.py
+++ b/osmosis_ai/rollout/eval/evaluation/cli.py
@@ -66,7 +66,7 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None:
"--dataset",
dest="dataset",
required=True,
- help="Path to dataset file (.json, .jsonl, or .parquet).",
+ help="Path to dataset file (.parquet recommended, .jsonl, or .csv).",
)
parser.add_argument(
diff --git a/osmosis_ai/rollout/eval/test_mode/cli.py b/osmosis_ai/rollout/eval/test_mode/cli.py
index 061dcedc..125751a2 100644
--- a/osmosis_ai/rollout/eval/test_mode/cli.py
+++ b/osmosis_ai/rollout/eval/test_mode/cli.py
@@ -85,7 +85,7 @@ def configure_parser(self, parser: argparse.ArgumentParser) -> None:
"--dataset",
dest="dataset",
required=True,
- help="Path to dataset file (.json, .jsonl, or .parquet).",
+ help="Path to dataset file (.parquet recommended, .jsonl, or .csv).",
)
parser.add_argument(
diff --git a/tests/unit/rollout/eval/common/test_dataset.py b/tests/unit/rollout/eval/common/test_dataset.py
index 0d2f7268..9dbf1386 100644
--- a/tests/unit/rollout/eval/common/test_dataset.py
+++ b/tests/unit/rollout/eval/common/test_dataset.py
@@ -30,30 +30,6 @@
class TestDatasetReader:
"""Tests for DatasetReader class."""
- def test_read_json_file(self, tmp_path: Path) -> None:
- """Test reading a valid JSON file."""
- data = [
- {
- "user_prompt": "What is 2+2?",
- "system_prompt": "You are a calculator.",
- "ground_truth": "4",
- },
- {
- "user_prompt": "What is 3+3?",
- "system_prompt": "You are a calculator.",
- "ground_truth": "6",
- },
- ]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
-
- reader = DatasetReader(str(file_path))
- rows = reader.read()
-
- assert len(rows) == 2
- assert rows[0]["user_prompt"] == "What is 2+2?"
- assert rows[1]["ground_truth"] == "6"
-
def test_read_jsonl_file(self, tmp_path: Path) -> None:
"""Test reading a valid JSONL file."""
lines = [
@@ -70,6 +46,79 @@ def test_read_jsonl_file(self, tmp_path: Path) -> None:
assert rows[0]["user_prompt"] == "Hello"
assert rows[1]["ground_truth"] == "Goodbye"
+ def test_read_csv_file(self, tmp_path: Path) -> None:
+ """Test reading a valid CSV file."""
+ content = (
+ "user_prompt,system_prompt,ground_truth\n"
+ "What is 2+2?,You are a calculator.,4\n"
+ "What is 3+3?,You are a calculator.,6\n"
+ )
+ file_path = tmp_path / "test.csv"
+ file_path.write_text(content)
+
+ reader = DatasetReader(str(file_path))
+ rows = reader.read()
+
+ assert len(rows) == 2
+ assert rows[0]["user_prompt"] == "What is 2+2?"
+ assert rows[1]["ground_truth"] == "6"
+
+ def test_csv_with_extra_columns(self, tmp_path: Path) -> None:
+ """Test that extra columns are preserved in CSV files."""
+ content = (
+ "user_prompt,system_prompt,ground_truth,difficulty,category\n"
+ "Question,System,Answer,easy,math\n"
+ )
+ file_path = tmp_path / "test.csv"
+ file_path.write_text(content)
+
+ reader = DatasetReader(str(file_path))
+ rows = reader.read()
+
+ assert len(rows) == 1
+ assert rows[0]["difficulty"] == "easy"
+ assert rows[0]["category"] == "math"
+
+ def test_csv_with_quoted_fields(self, tmp_path: Path) -> None:
+ """Test that CSV fields containing commas and quotes are parsed correctly."""
+ content = (
+ "system_prompt,user_prompt,ground_truth\n"
+ '"You are a helpful, friendly assistant.",What is 2 + 2?,4\n'
+ '"Say ""hello"".",Greet me,hello\n'
+ )
+ file_path = tmp_path / "test.csv"
+ file_path.write_text(content)
+
+ reader = DatasetReader(str(file_path))
+ rows = reader.read()
+
+ assert len(rows) == 2
+ assert rows[0]["system_prompt"] == "You are a helpful, friendly assistant."
+ assert rows[1]["system_prompt"] == 'Say "hello".'
+
+ def test_csv_missing_required_column(self, tmp_path: Path) -> None:
+ """Test that DatasetValidationError is raised for CSV with missing column."""
+ content = "user_prompt,system_prompt\nQuestion,System\n"
+ file_path = tmp_path / "test.csv"
+ file_path.write_text(content)
+
+ reader = DatasetReader(str(file_path))
+ with pytest.raises(DatasetValidationError) as exc_info:
+ reader.read()
+ assert "Missing required columns" in str(exc_info.value)
+ assert "ground_truth" in str(exc_info.value)
+
+ def test_csv_len(self, tmp_path: Path) -> None:
+ """Test that __len__ returns correct row count for CSV files."""
+ lines = ["user_prompt,system_prompt,ground_truth"]
+ for i in range(5):
+ lines.append(f"Q{i},sys,A{i}")
+ file_path = tmp_path / "test.csv"
+ file_path.write_text("\n".join(lines))
+
+ reader = DatasetReader(str(file_path))
+ assert len(reader) == 5
+
@pytest.mark.skipif(not HAS_PYARROW, reason="pyarrow not installed")
def test_read_parquet_file(self, tmp_path: Path) -> None:
"""Test reading a valid Parquet file."""
@@ -132,12 +181,18 @@ def test_parquet_len_uses_metadata(self, tmp_path: Path) -> None:
def test_read_with_limit(self, tmp_path: Path) -> None:
"""Test reading with limit parameter."""
- data = [
- {"user_prompt": f"Q{i}", "system_prompt": "sys", "ground_truth": f"A{i}"}
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": f"Q{i}",
+ "system_prompt": "sys",
+ "ground_truth": f"A{i}",
+ }
+ )
for i in range(10)
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
rows = reader.read(limit=3)
@@ -148,12 +203,18 @@ def test_read_with_limit(self, tmp_path: Path) -> None:
def test_read_with_offset(self, tmp_path: Path) -> None:
"""Test reading with offset parameter."""
- data = [
- {"user_prompt": f"Q{i}", "system_prompt": "sys", "ground_truth": f"A{i}"}
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": f"Q{i}",
+ "system_prompt": "sys",
+ "ground_truth": f"A{i}",
+ }
+ )
for i in range(10)
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
rows = reader.read(offset=5)
@@ -164,12 +225,18 @@ def test_read_with_offset(self, tmp_path: Path) -> None:
def test_read_with_limit_and_offset(self, tmp_path: Path) -> None:
"""Test reading with both limit and offset."""
- data = [
- {"user_prompt": f"Q{i}", "system_prompt": "sys", "ground_truth": f"A{i}"}
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": f"Q{i}",
+ "system_prompt": "sys",
+ "ground_truth": f"A{i}",
+ }
+ )
for i in range(10)
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
rows = reader.read(limit=3, offset=2)
@@ -180,15 +247,11 @@ def test_read_with_limit_and_offset(self, tmp_path: Path) -> None:
def test_case_insensitive_column_names(self, tmp_path: Path) -> None:
"""Test that column names are matched case-insensitively."""
- data = [
- {
- "USER_PROMPT": "Question",
- "System_Prompt": "System",
- "GROUND_TRUTH": "Answer",
- }
+ lines = [
+ '{"USER_PROMPT": "Question", "System_Prompt": "System", "GROUND_TRUTH": "Answer"}'
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
rows = reader.read()
@@ -201,18 +264,20 @@ def test_case_insensitive_column_names(self, tmp_path: Path) -> None:
def test_preserve_extra_columns(self, tmp_path: Path) -> None:
"""Test that extra columns are preserved in the result."""
- data = [
- {
- "user_prompt": "Question",
- "system_prompt": "System",
- "ground_truth": "Answer",
- "difficulty": "easy",
- "category": "math",
- "custom_field": 123,
- }
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": "Question",
+ "system_prompt": "System",
+ "ground_truth": "Answer",
+ "difficulty": "easy",
+ "category": "math",
+ "custom_field": 123,
+ }
+ )
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
rows = reader.read()
@@ -225,12 +290,18 @@ def test_preserve_extra_columns(self, tmp_path: Path) -> None:
def test_len_returns_row_count(self, tmp_path: Path) -> None:
"""Test that __len__ returns correct row count."""
- data = [
- {"user_prompt": f"Q{i}", "system_prompt": "sys", "ground_truth": f"A{i}"}
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": f"Q{i}",
+ "system_prompt": "sys",
+ "ground_truth": f"A{i}",
+ }
+ )
for i in range(5)
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
assert len(reader) == 5
@@ -238,20 +309,20 @@ def test_len_returns_row_count(self, tmp_path: Path) -> None:
def test_file_not_found(self) -> None:
"""Test that FileNotFoundError is raised for missing file."""
with pytest.raises(FileNotFoundError):
- DatasetReader("/nonexistent/path/data.json")
+ DatasetReader("/nonexistent/path/data.jsonl")
def test_unsupported_format(self, tmp_path: Path) -> None:
"""Test that DatasetParseError is raised for unsupported format."""
- file_path = tmp_path / "test.csv"
- file_path.write_text("a,b,c\n1,2,3")
+ file_path = tmp_path / "test.json"
+ file_path.write_text('[{"a": 1}]')
with pytest.raises(DatasetParseError) as exc_info:
DatasetReader(str(file_path))
assert "Unsupported file format" in str(exc_info.value)
- def test_invalid_json(self, tmp_path: Path) -> None:
- """Test that DatasetParseError is raised for invalid JSON."""
- file_path = tmp_path / "test.json"
+ def test_invalid_jsonl(self, tmp_path: Path) -> None:
+ """Test that DatasetParseError is raised for invalid JSONL."""
+ file_path = tmp_path / "test.jsonl"
file_path.write_text("not valid json")
reader = DatasetReader(str(file_path))
@@ -259,27 +330,19 @@ def test_invalid_json(self, tmp_path: Path) -> None:
reader.read()
assert "Invalid JSON" in str(exc_info.value)
- def test_json_not_array(self, tmp_path: Path) -> None:
- """Test that DatasetParseError is raised when JSON is not an array."""
- file_path = tmp_path / "test.json"
- file_path.write_text('{"key": "value"}')
-
- reader = DatasetReader(str(file_path))
- with pytest.raises(DatasetParseError) as exc_info:
- reader.read()
- assert "array of objects" in str(exc_info.value)
-
def test_missing_required_column(self, tmp_path: Path) -> None:
"""Test that DatasetValidationError is raised for missing column."""
- data = [
- {
- "user_prompt": "Question",
- "system_prompt": "System",
- # missing ground_truth
- }
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": "Question",
+ "system_prompt": "System",
+ # missing ground_truth
+ }
+ )
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
with pytest.raises(DatasetValidationError) as exc_info:
@@ -289,15 +352,17 @@ def test_missing_required_column(self, tmp_path: Path) -> None:
def test_null_value_rejected(self, tmp_path: Path) -> None:
"""Test that null values are rejected."""
- data = [
- {
- "user_prompt": None,
- "system_prompt": "System",
- "ground_truth": "Answer",
- }
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": None,
+ "system_prompt": "System",
+ "ground_truth": "Answer",
+ }
+ )
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
with pytest.raises(DatasetValidationError) as exc_info:
@@ -306,15 +371,17 @@ def test_null_value_rejected(self, tmp_path: Path) -> None:
def test_empty_string_rejected(self, tmp_path: Path) -> None:
"""Test that empty strings are rejected."""
- data = [
- {
- "user_prompt": " ", # whitespace only
- "system_prompt": "System",
- "ground_truth": "Answer",
- }
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": " ", # whitespace only
+ "system_prompt": "System",
+ "ground_truth": "Answer",
+ }
+ )
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
with pytest.raises(DatasetValidationError) as exc_info:
@@ -323,15 +390,17 @@ def test_empty_string_rejected(self, tmp_path: Path) -> None:
def test_non_string_value_rejected(self, tmp_path: Path) -> None:
"""Test that non-string values are rejected for required columns."""
- data = [
- {
- "user_prompt": 123, # number instead of string
- "system_prompt": "System",
- "ground_truth": "Answer",
- }
+ lines = [
+ json.dumps(
+ {
+ "user_prompt": 123, # number instead of string
+ "system_prompt": "System",
+ "ground_truth": "Answer",
+ }
+ )
]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
with pytest.raises(DatasetValidationError) as exc_info:
@@ -340,9 +409,9 @@ def test_non_string_value_rejected(self, tmp_path: Path) -> None:
def test_row_not_object_rejected(self, tmp_path: Path) -> None:
"""Test that non-object rows are rejected."""
- data = ["just a string", "another string"]
- file_path = tmp_path / "test.json"
- file_path.write_text(json.dumps(data))
+ lines = ['"just a string"', '"another string"']
+ file_path = tmp_path / "test.jsonl"
+ file_path.write_text("\n".join(lines))
reader = DatasetReader(str(file_path))
with pytest.raises(DatasetValidationError) as exc_info:
From 0c0d25e616247209f053ee7d86222c145f1021ec Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 13:46:13 -0800
Subject: [PATCH 35/60] Update README and CLI documentation for improved
clarity and structure
- Enhanced the README to clearly outline the two training modes: Local Rollout and Remote Rollout, including their respective workflows and example repositories.
- Updated CLI documentation to reflect new testing and evaluation commands, consolidating agent testing and evaluation sections for better usability.
- Removed outdated references and redundant sections to streamline content across documentation.
- Added links to example repositories for both training modes to facilitate user onboarding.
---
README.md | 124 +++++-----
docs/README.md | 43 ++--
docs/cli.md | 108 ++++-----
.../dataset-format.md => datasets.md} | 4 +-
docs/{rollout/eval.md => eval-mode.md} | 71 ++----
docs/local-rollout/mcp-tools.md | 139 ++++++++++++
docs/local-rollout/overview.md | 86 +++++++
docs/local-rollout/reward-functions.md | 122 ++++++++++
docs/local-rollout/reward-rubrics.md | 134 +++++++++++
.../agent-loop.md} | 8 +-
.../architecture.md | 6 +
.../{rollout => remote-rollout}/deployment.md | 0
docs/{rollout => remote-rollout}/examples.md | 10 +-
.../README.md => remote-rollout/overview.md} | 12 +-
docs/{rollout => remote-rollout}/testing.md | 2 +-
docs/rewards-api.md | 214 ++++++++++++++++++
docs/rewards.md | 189 ----------------
docs/{rollout => }/test-mode.md | 38 ++--
examples/README.md | 15 +-
19 files changed, 932 insertions(+), 393 deletions(-)
rename docs/{rollout/dataset-format.md => datasets.md} (83%)
rename docs/{rollout/eval.md => eval-mode.md} (91%)
create mode 100644 docs/local-rollout/mcp-tools.md
create mode 100644 docs/local-rollout/overview.md
create mode 100644 docs/local-rollout/reward-functions.md
create mode 100644 docs/local-rollout/reward-rubrics.md
rename docs/{rollout/api-reference.md => remote-rollout/agent-loop.md} (98%)
rename docs/{rollout => remote-rollout}/architecture.md (97%)
rename docs/{rollout => remote-rollout}/deployment.md (100%)
rename docs/{rollout => remote-rollout}/examples.md (98%)
rename docs/{rollout/README.md => remote-rollout/overview.md} (92%)
rename docs/{rollout => remote-rollout}/testing.md (98%)
create mode 100644 docs/rewards-api.md
delete mode 100644 docs/rewards.md
rename docs/{rollout => }/test-mode.md (91%)
diff --git a/README.md b/README.md
index aa25da2c..0b0677ac 100644
--- a/README.md
+++ b/README.md
@@ -1,85 +1,93 @@
# osmosis-ai
-A Python SDK for Osmosis LLM training workflows:
-- Reward/rubric validation helpers with strict type enforcement
-- Remote Rollout SDK for integrating agent frameworks with Osmosis training
-- Agent testing with external LLMs (`osmosis test`)
-- Model evaluation with custom eval functions and pass@k metrics (`osmosis eval`)
-- MCP tools support for git-sync workflows (`--mcp`)
+Python SDK for Osmosis AI training workflows. Supports two training modes with shared tooling for testing and evaluation.
+
+## Two Training Modes
+
+Osmosis supports **Local Rollout** and **Remote Rollout** as parallel approaches to training with reinforcement learning:
+
+| | Local Rollout | Remote Rollout |
+|--|--------------|----------------|
+| **How it works** | Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools via a GitHub-synced repo. | You implement and host a `RolloutAgentLoop` server. Full control over agent behavior. |
+| **Best for** | Standard tool-use agents, fast iteration, zero infrastructure | Custom agent architectures, complex orchestration, persistent environments |
+| **Example repo** | [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) | [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) |
## Installation
+Requires Python 3.10 or newer. For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
+
+### pip
+
```bash
-pip install osmosis-ai
+pip install osmosis-ai # Core SDK
+pip install osmosis-ai[server] # FastAPI server for Remote Rollout
+pip install osmosis-ai[mcp] # MCP tool support for Local Rollout
+pip install osmosis-ai[full] # All features
```
-Requires Python 3.10 or newer. For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
+### uv
-## Quick Start
+```bash
+uv add osmosis-ai # Core SDK
+uv add osmosis-ai[server] # FastAPI server for Remote Rollout
+uv add osmosis-ai[mcp] # MCP tool support for Local Rollout
+uv add osmosis-ai[full] # All features
+```
-```python
-from osmosis_ai import osmosis_reward
+## Local Rollout
-@osmosis_reward
-def simple_reward(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- """Basic exact match reward function."""
- return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
+Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools via a GitHub-synced repo. Best for standard tool-use agents, fast iteration, and zero infrastructure.
-# Use the reward function
-score = simple_reward("hello world", "hello world") # Returns 1.0
-```
+Get started with the example repo: **[osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)**
-```python
-from osmosis_ai import evaluate_rubric
+For details, see the [Local Rollout docs](docs/local-rollout/overview.md).
-solution = "The capital of France is Paris."
+## Remote Rollout
-# Export OPENAI_API_KEY in your shell before running this snippet.
-rubric_score = evaluate_rubric(
- rubric="Assistant must mention the verified capital city.",
- solution_str=solution,
- model_info={
- "provider": "openai",
- "model": "gpt-5",
- "api_key_env": "OPENAI_API_KEY",
- },
- ground_truth="Paris",
-)
+You implement and host a `RolloutAgentLoop` server. Full control over agent behavior, custom orchestration, and persistent environments.
-print(rubric_score) # -> 1.0 (full payload available via return_details=True)
-```
+Get started with the example repo: **[osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)**
+
+For details, see the [Remote Rollout docs](docs/remote-rollout/overview.md).
+
+## Testing & Evaluation
-## Remote Rollout SDK
+Both modes share the same `osmosis test` and `osmosis eval` CLI tools. See the example repos for usage:
-If you're integrating an agent loop with Osmosis remote rollout / TrainGate, see the [Rollout Quick Start](docs/rollout/README.md).
+- **Local Rollout**: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
+- **Remote Rollout**: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
+
+For CLI details, see [Test Mode](docs/test-mode.md), [Eval Mode](docs/eval-mode.md), and the [CLI Reference](docs/cli.md).
## Documentation
-| Topic | Description |
-|-------|-------------|
-| [Rewards & Rubrics](docs/rewards.md) | `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric` |
-| [CLI Reference](docs/cli.md) | All `osmosis` commands: auth, serve, test, eval, rubric tools |
-| **Remote Rollout SDK** | |
-| [Quick Start](docs/rollout/README.md) | Get a rollout server running in minutes |
-| [Architecture](docs/rollout/architecture.md) | Protocol design and agent lifecycle |
-| [API Reference](docs/rollout/api-reference.md) | Endpoints, schemas, types |
-| [Examples](docs/rollout/examples.md) | Agent implementations and utilities |
-| [Dataset Format](docs/rollout/dataset-format.md) | Supported formats and required columns |
-| [Test Mode](docs/rollout/test-mode.md) | Test agents with cloud LLMs |
-| [Evaluation](docs/rollout/eval.md) | Evaluate agents with eval functions and pass@k |
-| [Testing](docs/rollout/testing.md) | Unit tests and mock trainer |
-| [Deployment](docs/rollout/deployment.md) | Docker, health checks, production config |
+| Section | Topics |
+|---------|--------|
+| **Local Rollout** | |
+| [Overview](docs/local-rollout/overview.md) | When to choose, repo structure, setup |
+| [Reward Functions](docs/local-rollout/reward-functions.md) | `@osmosis_reward` decorator |
+| [Reward Rubrics](docs/local-rollout/reward-rubrics.md) | `@osmosis_rubric`, `evaluate_rubric` |
+| [MCP Tools](docs/local-rollout/mcp-tools.md) | `@mcp.tool()` definitions |
+| **Remote Rollout** | |
+| [Overview](docs/remote-rollout/overview.md) | Quick start, key concepts |
+| [Architecture](docs/remote-rollout/architecture.md) | Protocol design and lifecycle |
+| [Agent Loop Guide](docs/remote-rollout/agent-loop.md) | API reference |
+| [Examples](docs/remote-rollout/examples.md) | Agent implementations |
+| [Testing](docs/remote-rollout/testing.md) | Unit tests and mock trainer |
+| [Deployment](docs/remote-rollout/deployment.md) | Docker, production config |
+| **Shared** | |
+| [Dataset Format](docs/datasets.md) | Parquet, JSONL, CSV formats |
+| [Test Mode](docs/test-mode.md) | `osmosis test` |
+| [Eval Mode](docs/eval-mode.md) | `osmosis eval` with pass@k |
+| **Reference** | |
+| [Rewards API](docs/rewards-api.md) | Decorator and function signatures |
+| [CLI Reference](docs/cli.md) | All `osmosis` commands |
Full documentation index: [docs/README.md](docs/README.md)
-## Running Examples
-
-```bash
-PYTHONPATH=. python examples/reward_functions.py
-PYTHONPATH=. python examples/rubric_functions.py # Uncomment the provider you need before running
-```
+## Examples
-See the [`examples/`](examples/) directory for all examples.
+The [`examples/`](examples/) directory contains standalone SDK API examples (reward functions, rubric evaluation). See [`examples/README.md`](examples/README.md) for details.
## Contributing
@@ -93,3 +101,5 @@ MIT License - see [LICENSE](LICENSE) file for details.
- [Homepage](https://github.com/Osmosis-AI/osmosis-sdk-python)
- [Issues](https://github.com/Osmosis-AI/osmosis-sdk-python/issues)
+- [Local Rollout Example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
+- [Remote Rollout Example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
diff --git a/docs/README.md b/docs/README.md
index bcb9cbba..8a4cc1a6 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,34 +1,47 @@
# Osmosis SDK Documentation
-Python SDK for Osmosis AI training workflows, providing reward function validation, rubric evaluation via LLM-as-judge, a Remote Rollout SDK for integrating agent frameworks with Osmosis training, and CLI tools for testing, evaluation, and workspace management.
+Python SDK for Osmosis AI training workflows. The SDK supports two parallel training modes -- **Local Rollout** and **Remote Rollout** -- along with shared tools for testing, evaluation, and reward scoring.
## Getting Started
- [Installation & Quick Start](../README.md)
- [Contributing](../CONTRIBUTING.md)
-## Rewards & Rubrics
+## Local Rollout
-- [Rewards & Rubrics Guide](./rewards.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
-- [Example Code](../examples/) -- reward functions, rubric configs, sample data
+Git-sync mode where Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools.
+
+- [Overview](./local-rollout/overview.md) -- what is local rollout, when to choose it, repository structure
+- [Reward Functions](./local-rollout/reward-functions.md) -- `@osmosis_reward` decorator usage and examples
+- [Reward Rubrics](./local-rollout/reward-rubrics.md) -- `@osmosis_rubric` and `evaluate_rubric` usage
+- [MCP Tools](./local-rollout/mcp-tools.md) -- defining `@mcp.tool()` functions for the agent
+
+## Remote Rollout
+
+Self-hosted mode where you implement and run the agent loop as a server.
-## Remote Rollout SDK
+- [Overview](./remote-rollout/overview.md) -- quick start and key concepts
+- [Architecture](./remote-rollout/architecture.md) -- protocol design and agent lifecycle
+- [Agent Loop Guide](./remote-rollout/agent-loop.md) -- API reference for classes, schemas, types
+- [Examples](./remote-rollout/examples.md) -- agent implementations and utilities
+- [Testing](./remote-rollout/testing.md) -- unit tests and mock trainer
+- [Deployment](./remote-rollout/deployment.md) -- Docker, health checks, production config
-- [Quick Start](./rollout/README.md) -- get a rollout server running in minutes
-- [Architecture](./rollout/architecture.md) -- protocol design and agent lifecycle
-- [API Reference](./rollout/api-reference.md) -- endpoints, schemas, types
-- [Examples](./rollout/examples.md) -- agent implementations and utilities
-- [Dataset Format](./rollout/dataset-format.md) -- supported formats and required columns
-- [Test Mode](./rollout/test-mode.md) -- test agents with cloud LLMs
-- [Evaluation](./rollout/eval.md) -- evaluate agents against datasets
-- [Testing](./rollout/testing.md) -- unit tests and mock trainer
-- [Deployment](./rollout/deployment.md) -- Docker, health checks, production config
+## Shared Concepts
-## CLI Reference
+Used by both Local Rollout and Remote Rollout modes.
+- [Dataset Format](./datasets.md) -- supported formats and required columns
+- [Test Mode](./test-mode.md) -- test agents with cloud LLMs (`osmosis test`)
+- [Eval Mode](./eval-mode.md) -- evaluate agents with eval functions and pass@k (`osmosis eval`)
+
+## Reference
+
+- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric` API details
- [CLI Reference](./cli.md) -- all `osmosis` commands
## Other
+- [Example Code](../examples/) -- reward functions, rubric configs, sample data
- [Security Policy](../SECURITY.md)
- [License](../LICENSE)
diff --git a/docs/cli.md b/docs/cli.md
index 3b78040d..f313c350 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -1,6 +1,6 @@
# CLI Reference
-Installing the SDK provides a lightweight CLI available as `osmosis` (aliases: `osmosis_ai`, `osmosis-ai`). The CLI covers authentication, workspace management, rollout server operation, agent testing and evaluation, and rubric tools. It automatically loads `.env` from the current working directory via `python-dotenv`.
+Installing the SDK provides a lightweight CLI available as `osmosis` (aliases: `osmosis_ai`, `osmosis-ai`). The CLI automatically loads `.env` from the current working directory via `python-dotenv`.
## Authentication
@@ -44,11 +44,9 @@ Display the current user and all workspaces:
osmosis whoami
```
-## Workspace Management
-
### osmosis workspace
-Manage multiple workspaces after logging in. You can log in to multiple workspaces and switch between them. Each workspace maintains its own credentials and role information.
+Manage multiple workspaces after logging in. You can log in to multiple workspaces and switch between them.
```bash
# List all logged-in workspaces
@@ -61,6 +59,53 @@ osmosis workspace current
osmosis workspace switch
```
+## Testing Your Agent
+
+### osmosis test
+
+Test your agent locally against a dataset using external LLMs. Works with both Local Rollout (MCP tools) and Remote Rollout (RolloutAgentLoop) agents.
+
+```bash
+# Remote Rollout: test an AgentLoop implementation
+osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini
+
+# Local Rollout: test MCP tools directly (no AgentLoop needed)
+osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
+
+# Interactive step-by-step debugging
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
+
+# Interactive with MCP tools
+osmosis test --mcp ./mcp -d data.jsonl --interactive
+```
+
+See [Test Mode](./test-mode.md) for full documentation on dataset format, interactive mode, and all options.
+
+## Evaluating Models
+
+### osmosis eval
+
+Evaluate trained models with custom eval functions and pass@k metrics. Works with both Local Rollout (MCP tools) and Remote Rollout (RolloutAgentLoop) agents.
+
+```bash
+# Remote Rollout: benchmark a trained model at a serving endpoint
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
+ --eval-fn rewards:compute_reward \
+ --model my-finetuned-model --base-url http://localhost:8000/v1
+
+# Local Rollout: evaluate MCP tools directly
+osmosis eval --mcp ./mcp -d data.jsonl \
+ --eval-fn rewards:compute_reward \
+ --model openai/gpt-5-mini
+
+# pass@k with 5 runs per row
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
+ --eval-fn rewards:compute_reward --n 5 \
+ --model my-finetuned-model --base-url http://localhost:8000/v1
+```
+
+See [Eval Mode](./eval-mode.md) for full documentation on eval functions, pass@k metrics, and output formats.
+
## Remote Rollout Server
### osmosis serve
@@ -99,50 +144,6 @@ Validate an agent loop before starting the server (checks tools, async run metho
osmosis validate -m my_agent:agent_loop
```
-## Agent Testing
-
-### osmosis test
-
-Test your agent loop locally against a dataset using external LLMs:
-
-```bash
-# Batch test with OpenAI
-osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini
-
-# Interactive step-by-step debugging
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
-
-# Git-sync users: test MCP tools directly (no AgentLoop code needed)
-osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-```
-
-See [Test Mode](./rollout/test-mode.md) for full documentation on dataset format, interactive mode, and MCP tool testing.
-
-## Agent Evaluation
-
-### osmosis eval
-
-Evaluate trained models with custom eval functions and pass@k metrics:
-
-```bash
-# Benchmark a trained model at a serving endpoint
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# pass@k with 5 runs per row
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Git-sync users: evaluate MCP tools directly
-osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model openai/gpt-5-mini
-```
-
-See [Evaluation](./rollout/eval.md) for full documentation on eval functions, pass@k metrics, and output formats.
-
## Rubric Tools
### osmosis preview
@@ -171,7 +172,7 @@ osmosis eval-rubric --rubric support_followup --data examples/sample_data.jsonl
**Command split** (development-stage breaking change):
- `osmosis eval-rubric` evaluates JSONL conversations against hosted rubrics.
-- `osmosis eval` runs rollout eval functions against `RolloutAgentLoop` datasets.
+- `osmosis eval` runs rollout eval functions against agent datasets.
**Options:**
@@ -188,10 +189,9 @@ osmosis eval-rubric --rubric support_followup --data examples/sample_data.jsonl
- When delegating to a custom `@osmosis_rubric` function, the CLI enriches `extra_info` with the active `provider`, `model`, `rubric`, score bounds, any configured `system_prompt`, the resolved `original_input`, and the record's metadata/extra fields so the decorator's required entries are always present.
- Rubric configuration files intentionally reject `extra_info`; provide per-example context through the dataset instead.
-See [Rewards & Rubrics](./rewards.md) for details on `@osmosis_reward`, `@osmosis_rubric`, and `evaluate_rubric`.
-
## See Also
-- [Rewards & Rubrics](./rewards.md)
-- [Rollout Test Mode](./rollout/test-mode.md)
-- [Rollout Evaluation](./rollout/eval.md)
+- [Test Mode](./test-mode.md) -- full `osmosis test` documentation
+- [Eval Mode](./eval-mode.md) -- full `osmosis eval` documentation
+- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
+- [Dataset Format](./datasets.md) -- supported formats and required columns
diff --git a/docs/rollout/dataset-format.md b/docs/datasets.md
similarity index 83%
rename from docs/rollout/dataset-format.md
rename to docs/datasets.md
index d6f3ade0..db4c6028 100644
--- a/docs/rollout/dataset-format.md
+++ b/docs/datasets.md
@@ -1,6 +1,6 @@
# Dataset Format
-Datasets define the test cases for evaluating and testing your rollout agents. They provide the prompts, expected outputs, and optional metadata used by both [test mode](./test-mode.md) and [evaluation](./eval.md).
+Datasets define the test cases for testing and evaluating your agents. They provide the prompts, expected outputs, and optional metadata used by both [test mode](./test-mode.md) and [eval mode](./eval-mode.md). Datasets are used by both **Local Rollout** and **Remote Rollout** modes.
## Supported Formats
@@ -56,4 +56,4 @@ Any columns beyond the three required ones are passed to `RolloutRequest.metadat
## See Also
- [Test Mode](./test-mode.md) -- test agents locally with external LLMs
-- [Evaluation](./eval.md) -- evaluate agents with eval functions and pass@k
+- [Eval Mode](./eval-mode.md) -- evaluate agents with eval functions and pass@k
diff --git a/docs/rollout/eval.md b/docs/eval-mode.md
similarity index 91%
rename from docs/rollout/eval.md
rename to docs/eval-mode.md
index 2cb7de30..101f0b4c 100644
--- a/docs/rollout/eval.md
+++ b/docs/eval-mode.md
@@ -1,6 +1,6 @@
# Eval Mode
-Evaluate your trained models by running agent implementations against datasets with custom eval functions, statistical analysis, and pass@k metrics.
+Evaluate your trained models by running agent implementations against datasets with custom eval functions, statistical analysis, and pass@k metrics. Works with both **Local Rollout** (MCP tools) and **Remote Rollout** (RolloutAgentLoop) agents.
## Overview
@@ -9,7 +9,7 @@ Eval mode is designed for **evaluating trained models**. Connect to any OpenAI-c
Key capabilities:
- **Benchmark trained models** against eval functions by connecting to serving endpoints
-- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp` (for git-sync users)
+- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp`
- Run multiple trials per row for pass@k analysis
- Use existing `@osmosis_reward` functions or full-context eval functions
- Get statistical summaries (mean, std, min, max) per eval function
@@ -22,35 +22,28 @@ Key capabilities:
## Quick Start
-### Benchmarking a Trained Model
+### Benchmarking with Remote Rollout Agent
The primary use case is connecting to a model serving endpoint:
```bash
# Benchmark a trained model served at an endpoint
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model \
--base-url http://localhost:8000/v1
-# Benchmark against osmosis-serving
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-trained-model \
- --base-url https://inference.osmosis.ai/v1 \
- --api-key $OSMOSIS_API_KEY
-
# Multiple eval functions
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:exact_match \
--eval-fn rewards:partial_match \
--model my-finetuned-model \
--base-url http://localhost:8000/v1
```
-### Git-Sync Users (MCP Tools)
+### Benchmarking with Local Rollout MCP Tools
-If you use the **git-sync** workflow and provide MCP tools (`@mcp.tool()`) instead of a `RolloutAgentLoop`, use `--mcp` to point at your MCP directory. The SDK automatically loads all registered tools and runs a standard agent loop — no AgentLoop code required.
+If you use the **git-sync** workflow and provide MCP tools (`@mcp.tool()`) instead of a `RolloutAgentLoop`, use `--mcp` to point at your MCP directory. The SDK automatically loads all registered tools and runs a standard agent loop.
```bash
# Install MCP support
@@ -97,19 +90,19 @@ Compare your trained model against a baseline model. The SDK runs both models on
```bash
# Compare trained model vs GPT-5-mini baseline
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1 \
--baseline-model openai/gpt-5-mini
# Compare two serving endpoints
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-model-v2 --base-url http://localhost:8000/v1 \
--baseline-model my-model-v1 --baseline-base-url http://localhost:8001/v1
# Baseline with explicit API key
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1 \
--baseline-model anthropic/claude-sonnet-4-5 --baseline-api-key $ANTHROPIC_API_KEY
@@ -125,13 +118,13 @@ Use `--batch-size` to run multiple requests in parallel:
```bash
# Run 5 concurrent requests
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1 \
--batch-size 5
# Combine with pass@k — 10 runs per row, 5 concurrent
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward --n 10 --batch-size 5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
```
@@ -140,12 +133,12 @@ osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
```bash
# pass@k with 5 runs per row
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward --n 5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
# Custom pass threshold (default is 1.0)
-osmosis eval -m my_agent:MyAgentLoop -d data.jsonl \
+osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward --n 5 --pass-threshold 0.5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
```
@@ -226,7 +219,7 @@ The `--model` parameter should match the model name as registered in the serving
## Dataset Format
-See [Dataset Format](./dataset-format.md) for supported formats and required columns.
+See [Dataset Format](./datasets.md) for supported formats and required columns.
---
@@ -319,8 +312,8 @@ osmosis eval [OPTIONS]
| Option | Description |
|--------|-------------|
-| `-m, --module MODULE` | Module path to agent loop (format: `module:attribute`). Use for remote-rollout users who implement `RolloutAgentLoop`. |
-| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for git-sync users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
+| `-m, --module MODULE` | Module path to agent loop (format: `module:attribute`). Use for Remote Rollout users who implement `RolloutAgentLoop`. |
+| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for Local Rollout (git-sync) users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
### Model Options
@@ -362,17 +355,17 @@ osmosis eval [OPTIONS]
### Examples
```bash
-# Benchmark trained model at an endpoint
+# Remote Rollout: benchmark trained model at an endpoint
osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1
-# Git-sync: evaluate MCP tools without writing an AgentLoop
+# Local Rollout: evaluate MCP tools
osmosis eval --mcp ./mcp -d data.jsonl \
--eval-fn reward_fn:compute_reward \
--model openai/gpt-5-mini
-# Git-sync: MCP tools with trained model endpoint
+# Local Rollout: MCP tools with trained model endpoint
osmosis eval --mcp ./mcp -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1
@@ -388,32 +381,16 @@ osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward --n 5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
-# Lenient pass threshold
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 --pass-threshold 0.5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Baseline comparison: trained model vs external LLM
+# Baseline comparison
osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1 \
--baseline-model openai/gpt-5-mini -o results.json
-# Baseline comparison: two serving endpoints
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-model-v2 --base-url http://localhost:8000/v1 \
- --baseline-model my-model-v1 --baseline-base-url http://localhost:8001/v1
-
# Concurrent execution (5 runs at a time)
osmosis eval -m my_agent:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward --batch-size 5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Benchmark subset
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --limit 10 --offset 50 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
```
---
@@ -742,7 +719,7 @@ See [Test Mode - Environment Variables](./test-mode.md#environment-variables) fo
## See Also
- [Test Mode](./test-mode.md) - Test agent logic with external LLMs
-- [Architecture](./architecture.md) - System design overview
-- [API Reference](./api-reference.md) - Complete SDK API documentation
-- [Examples](./examples.md) - Working code examples
+- [Dataset Format](./datasets.md) - Supported formats and required columns
+- [Remote Rollout Examples](./remote-rollout/examples.md) - Working code examples
+- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) - MCP tool definition
- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
diff --git a/docs/local-rollout/mcp-tools.md b/docs/local-rollout/mcp-tools.md
new file mode 100644
index 00000000..784ea385
--- /dev/null
+++ b/docs/local-rollout/mcp-tools.md
@@ -0,0 +1,139 @@
+# MCP Tools
+
+In Local Rollout mode, the agent's callable tools are defined as [MCP](https://modelcontextprotocol.io/) (Model Context Protocol) tool functions using [FastMCP](https://github.com/jlowin/fastmcp). During training, Osmosis loads your tools and makes them available to the LLM in the agent loop.
+
+## Defining Tools
+
+Tools are Python functions decorated with `@mcp.tool()`. Type hints and docstrings are used to generate the tool schema automatically.
+
+```python
+# mcp/tools/math.py
+from server import mcp
+
+@mcp.tool()
+def multiply(first_val: float, second_val: float) -> float:
+ """Calculate the product of two numbers.
+
+ Args:
+ first_val: the first value to be multiplied
+ second_val: the second value to be multiplied
+ """
+ return round(first_val * second_val, 4)
+```
+
+The tool name, description, and parameter schemas are derived from the function name, docstring, and type annotations. Keep docstrings clear and descriptive -- the LLM uses them to decide when and how to call your tools.
+
+## Folder Structure
+
+The `mcp/` directory must contain a `main.py` entry point that creates a `FastMCP` instance and imports all tool modules:
+
+```
+mcp/
+├── main.py # Entry point
+├── server/
+│ ├── __init__.py
+│ └── mcp_server.py # Creates the FastMCP instance
+└── tools/
+ ├── __init__.py # Imports all tool modules
+ ├── math.py # @mcp.tool() functions
+ ├── api_helpers.py
+ └── ml_utils.py
+```
+
+### main.py
+
+```python
+# mcp/main.py
+from server import mcp
+from tools import * # Triggers @mcp.tool() registrations
+
+if __name__ == "__main__":
+ import argparse
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--host', default='0.0.0.0')
+ parser.add_argument('--port', type=int, default=8080)
+ args = parser.parse_args()
+ mcp.run(transport="http", host=args.host, port=args.port)
+```
+
+### server/mcp_server.py
+
+```python
+# mcp/server/mcp_server.py
+from fastmcp import FastMCP
+
+mcp = FastMCP("OsmosisTools")
+```
+
+### tools/__init__.py
+
+```python
+# mcp/tools/__init__.py
+from .math import *
+from .api_helpers import *
+```
+
+## Simpler Structure
+
+For small projects, you can put everything in a single `main.py`:
+
+```python
+# mcp/main.py
+from fastmcp import FastMCP
+
+mcp = FastMCP("my_tools")
+
+@mcp.tool()
+def add(a: float, b: float) -> str:
+ """Add two numbers."""
+ return str(a + b)
+
+@mcp.tool()
+def lookup(key: str) -> str:
+ """Look up a value by key."""
+ data = {"pi": "3.14159", "e": "2.71828"}
+ return data.get(key, "not found")
+```
+
+## Testing MCP Tools Locally
+
+Use `osmosis test` with the `--mcp` flag to test your tools against a dataset:
+
+```bash
+# Install MCP support
+pip install "osmosis-ai[mcp]"
+
+# Batch test
+osmosis test --mcp ./mcp -d test_data.jsonl --model openai/gpt-5-mini
+
+# Interactive debugging -- step through each LLM call
+osmosis test --mcp ./mcp -d test_data.jsonl --interactive
+
+# Evaluate with a reward function
+osmosis eval --mcp ./mcp -d test_data.jsonl \
+ --eval-fn reward_fn.compute_reward:numbers_match_reward \
+ --model openai/gpt-5-mini
+```
+
+### How `--mcp` Works
+
+When you pass `--mcp ./mcp`, the SDK:
+
+1. Imports `mcp/main.py`, which triggers all `@mcp.tool()` registrations
+2. Discovers registered tools and converts them to OpenAI function-calling schemas
+3. Runs a built-in agent loop that calls the LLM, executes tool calls against your MCP functions, and repeats until the LLM stops calling tools or `--max-turns` is reached
+
+This means you can iterate on tools locally, then push to GitHub for Osmosis to sync -- no `RolloutAgentLoop` code needed.
+
+> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use `--mcp` for Local Rollout projects; use `-m` for Remote Rollout projects that implement `RolloutAgentLoop`.
+
+## Example Repository
+
+See the complete working example with MCP tools, reward functions, and rubrics: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
+
+## See Also
+
+- [Test Mode](../test-mode.md) -- full documentation for `osmosis test`
+- [Eval Mode](../eval-mode.md) -- full documentation for `osmosis eval`
+- [Local Rollout Overview](./overview.md) -- repository structure and setup
+- [Dataset Format](../datasets.md) -- dataset format for testing and evaluation
diff --git a/docs/local-rollout/overview.md b/docs/local-rollout/overview.md
new file mode 100644
index 00000000..c5c05f0b
--- /dev/null
+++ b/docs/local-rollout/overview.md
@@ -0,0 +1,86 @@
+# Local Rollout
+
+Local Rollout is one of two training modes supported by the Osmosis platform. In this mode, Osmosis manages the entire agent loop -- you provide **reward functions**, **rubric evaluators**, and optionally **MCP tools** via a GitHub-synced repository. The training infrastructure handles LLM inference, tool execution, and trajectory collection automatically.
+
+## When to Choose Local Rollout
+
+| Consideration | Local Rollout | Remote Rollout |
+|---------------|--------------|----------------|
+| **Agent logic** | Managed by Osmosis | You implement `RolloutAgentLoop` |
+| **Tool execution** | MCP tools synced from GitHub | Custom tool code in your server |
+| **Infrastructure** | None -- Osmosis handles everything | You host a RolloutServer |
+| **Flexibility** | Standard agent loop, configurable tools | Full control over agent behavior |
+| **Best for** | Tool-use tasks with standard agent patterns | Custom agent architectures, complex orchestration |
+
+Choose **Local Rollout** when:
+- Your task follows a standard tool-use agent pattern (LLM calls tools in a loop)
+- You want zero infrastructure to manage
+- Your tools can be expressed as MCP `@mcp.tool()` functions
+- You want fast iteration via git push
+
+Choose **Remote Rollout** when you need full control over the agent loop, custom orchestration logic, or tools that require a persistent server environment. See the [Remote Rollout overview](../remote-rollout/overview.md).
+
+## How It Works
+
+1. **Connect your GitHub repo** to Osmosis via the platform UI
+2. Osmosis syncs your repo and discovers:
+ - `reward_fn/` -- `@osmosis_reward` functions for deterministic scoring
+ - `reward_rubric/` -- `@osmosis_rubric` functions for LLM-as-judge scoring
+ - `mcp/` -- `@mcp.tool()` functions the agent can call during rollouts
+3. **Create a training run** on the platform, selecting "Local Rollout" mode
+4. Osmosis runs the agent loop on its training cluster, calling your MCP tools and scoring outputs with your reward/rubric functions
+
+## Required Repository Structure
+
+```
+your-repo/
+├── mcp/ # MCP tools (agent's callable functions)
+│ ├── main.py # Entry point -- creates FastMCP instance
+│ ├── server/ # Server setup
+│ │ └── mcp_server.py
+│ └── tools/ # Tool implementations
+│ ├── __init__.py
+│ └── math.py
+├── reward_fn/ # Reward functions (deterministic scoring)
+│ └── compute_reward.py
+├── reward_rubric/ # Rubric evaluators (LLM-as-judge scoring)
+│ ├── reward_rubric_openai.py
+│ └── reward_rubric_anthropic.py
+├── test_data.jsonl # Sample dataset for local testing
+└── pyproject.toml
+```
+
+All three directories (`mcp/`, `reward_fn/`, `reward_rubric/`) are optional -- include only what your training run needs.
+
+## Local Testing
+
+Before pushing to GitHub, test your setup locally:
+
+```bash
+# Install MCP support
+pip install "osmosis-ai[mcp]"
+
+# Test MCP tools against a dataset
+osmosis test --mcp ./mcp -d test_data.jsonl --model openai/gpt-5-mini
+
+# Interactive debugging
+osmosis test --mcp ./mcp -d test_data.jsonl --interactive
+
+# Evaluate with your reward function
+osmosis eval --mcp ./mcp -d test_data.jsonl \
+ --eval-fn reward_fn.compute_reward:numbers_match_reward \
+ --model openai/gpt-5-mini
+```
+
+See [Test Mode](../test-mode.md) and [Eval Mode](../eval-mode.md) for full documentation.
+
+## Example Repository
+
+See the complete working example: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
+
+## Next Steps
+
+- [Reward Functions](./reward-functions.md) -- define `@osmosis_reward` scoring functions
+- [Reward Rubrics](./reward-rubrics.md) -- define `@osmosis_rubric` LLM-as-judge evaluators
+- [MCP Tools](./mcp-tools.md) -- define `@mcp.tool()` functions for the agent
+- [Dataset Format](../datasets.md) -- dataset format for testing and evaluation
diff --git a/docs/local-rollout/reward-functions.md b/docs/local-rollout/reward-functions.md
new file mode 100644
index 00000000..d2c5eb63
--- /dev/null
+++ b/docs/local-rollout/reward-functions.md
@@ -0,0 +1,122 @@
+# Reward Functions
+
+Reward functions provide deterministic, code-based scoring for LLM outputs during training. They are used in Local Rollout mode to drive reinforcement learning -- the training loop calls your reward function after each rollout to compute a scalar score.
+
+## @osmosis_reward
+
+All functions decorated with `@osmosis_reward` must have exactly this signature:
+
+```python
+from osmosis_ai import osmosis_reward
+
+@osmosis_reward
+def your_function(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ # Your reward logic here
+ return float_score
+```
+
+### Parameters
+
+- **`solution_str: str`** -- The solution string to evaluate (required).
+- **`ground_truth: str`** -- The correct/expected answer (required).
+- **`extra_info: dict = None`** -- Optional dictionary for additional configuration.
+- **`**kwargs`** -- Required for platform compatibility.
+
+### Return Value
+
+- **`-> float`** -- Must return a float value representing the reward score.
+
+The decorator raises a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
+
+## Examples
+
+### Exact Match
+
+```python
+@osmosis_reward
+def exact_match(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ """Simple exact string matching."""
+ return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
+```
+
+### Case-Insensitive Match with Partial Credit
+
+```python
+@osmosis_reward
+def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ """Case-insensitive string matching with partial credit."""
+ match = solution_str.lower().strip() == ground_truth.lower().strip()
+
+ if extra_info and 'partial_credit' in extra_info:
+ if not match and extra_info['partial_credit']:
+ len_diff = abs(len(solution_str) - len(ground_truth))
+ if len_diff <= 2:
+ return 0.5
+
+ return 1.0 if match else 0.0
+```
+
+### Numeric Tolerance
+
+```python
+@osmosis_reward
+def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ """Numeric comparison with configurable tolerance."""
+ try:
+ solution_num = float(solution_str.strip())
+ truth_num = float(ground_truth.strip())
+
+ tolerance = extra_info.get('tolerance', 0.01) if extra_info else 0.01
+ return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
+ except ValueError:
+ return 0.0
+```
+
+### Extracting Solutions from Agent Output
+
+A common pattern for tool-use agents is extracting a final numeric answer from a markdown-formatted response:
+
+```python
+import re
+from osmosis_ai import osmosis_reward
+
+def extract_solution(solution_str):
+ """Extract the first number after a #### heading."""
+ solution = re.search(r'####\s*([-+]?\d*\.?\d+)', solution_str)
+ if not solution:
+ return None
+ return solution.group(1)
+
+@osmosis_reward
+def numbers_match_reward(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ extracted = extract_solution(solution_str)
+ try:
+ sol_val = float(extracted)
+ except (TypeError, ValueError):
+ return 0.0
+
+ gt_val = float(ground_truth)
+ if abs(gt_val - sol_val) < 1e-7:
+ return 1.0
+ return 0.0
+```
+
+This example is from the [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) repository (`reward_fn/compute_reward.py`).
+
+## File Placement
+
+In a Local Rollout repository, place reward functions in the `reward_fn/` directory:
+
+```
+reward_fn/
+└── compute_reward.py # Contains @osmosis_reward functions
+```
+
+Osmosis discovers and syncs all `@osmosis_reward`-decorated functions from this directory.
+
+## See Also
+
+- [Reward Rubrics](./reward-rubrics.md) -- LLM-as-judge scoring with `@osmosis_rubric`
+- [Rewards API Reference](../rewards-api.md) -- full API reference for decorators and `evaluate_rubric`
+- [Local Rollout Overview](./overview.md) -- repository structure and setup
+- [Example Code](../../examples/) -- reward functions, rubric configs, sample data
diff --git a/docs/local-rollout/reward-rubrics.md b/docs/local-rollout/reward-rubrics.md
new file mode 100644
index 00000000..83b85054
--- /dev/null
+++ b/docs/local-rollout/reward-rubrics.md
@@ -0,0 +1,134 @@
+# Reward Rubrics
+
+Rubric evaluations use an external LLM as a judge to score agent outputs against natural-language criteria. This is useful when correctness is subjective or hard to express as a deterministic function.
+
+## @osmosis_rubric
+
+Rubric functions decorated with `@osmosis_rubric` must match this signature:
+
+```python
+from osmosis_ai import osmosis_rubric
+
+@osmosis_rubric
+def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+ # Your rubric logic here
+ return float_score
+```
+
+> The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
+
+### Required `extra_info` Fields
+
+- **`provider`** -- Non-empty string identifying the judge provider.
+- **`model`** -- Non-empty string naming the provider model to call.
+- **`rubric`** -- Natural-language rubric instructions for the judge model.
+- **`api_key` / `api_key_env`** -- Supply either the raw key or the environment variable name that exposes it.
+
+### Optional `extra_info` Fields
+
+- **`system_prompt`** -- Optional string prepended to the provider's base system prompt.
+- **`score_min` / `score_max`** -- Optional numeric overrides for the expected score range.
+- **`model_info_overrides`** -- Optional dict merged into the provider configuration.
+
+Additional keys are passthrough and can be used for custom configuration. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
+
+> Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
+
+## evaluate_rubric
+
+`evaluate_rubric` talks to hosted LLM providers through [LiteLLM](https://github.com/BerriAI/litellm), a unified interface supporting 100+ providers. Every provider returns a strict JSON object with `{"score": number, "explanation": string}`.
+
+### Basic Usage
+
+```python
+from osmosis_ai import evaluate_rubric
+
+solution = "The capital of France is Paris."
+
+rubric_score = evaluate_rubric(
+ rubric="Assistant must mention the verified capital city.",
+ solution_str=solution,
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5",
+ "api_key_env": "OPENAI_API_KEY",
+ },
+ ground_truth="Paris",
+)
+
+print(rubric_score) # -> 1.0 (full payload available via return_details=True)
+```
+
+### Credentials
+
+Credentials are resolved from environment variables by default:
+
+- `OPENAI_API_KEY` for OpenAI
+- `ANTHROPIC_API_KEY` for Anthropic
+- `GEMINI_API_KEY` for Google Gemini
+- `XAI_API_KEY` for xAI
+- `OPENROUTER_API_KEY` for OpenRouter
+- `CEREBRAS_API_KEY` for Cerebras
+
+Override the environment variable name with `model_info={"api_key_env": "CUSTOM_ENV_NAME"}` when needed, or supply an inline secret with `model_info={"api_key": "sk-..."}` for ephemeral credentials.
+
+`api_key` and `api_key_env` are mutually exclusive. When `api_key` is present and non-empty it is used directly, skipping any environment lookup. Otherwise the resolver falls back to `api_key_env` (or the provider default) and pulls the value from your local environment with `os.getenv`.
+
+## Provider Examples
+
+### OpenAI
+
+```python
+from osmosis_ai import osmosis_rubric, evaluate_rubric
+
+@osmosis_rubric
+def compute_rubric_score_openai(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+ return evaluate_rubric(
+ rubric="Evaluate whether the solution correctly matches the expected answer.",
+ solution_str=solution_str,
+ ground_truth=ground_truth,
+ model_info={
+ "provider": "openai",
+ "model": "gpt-5-mini",
+ "api_key_env": "OPENAI_API_KEY",
+ },
+ )
+```
+
+### Anthropic
+
+```python
+@osmosis_rubric
+def compute_rubric_score_anthropic(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+ return evaluate_rubric(
+ rubric="Evaluate whether the solution correctly matches the expected answer.",
+ solution_str=solution_str,
+ ground_truth=ground_truth,
+ model_info={
+ "provider": "anthropic",
+ "model": "claude-sonnet-4-5-20250929",
+ "api_key_env": "ANTHROPIC_API_KEY",
+ },
+ )
+```
+
+Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
+
+## File Placement
+
+In a Local Rollout repository, place rubric functions in the `reward_rubric/` directory:
+
+```
+reward_rubric/
+├── reward_rubric_openai.py
+├── reward_rubric_anthropic.py
+└── reward_rubric_xai.py
+```
+
+Osmosis discovers and syncs all `@osmosis_rubric`-decorated functions from this directory. See the [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) for working examples.
+
+## See Also
+
+- [Reward Functions](./reward-functions.md) -- deterministic scoring with `@osmosis_reward`
+- [Rewards API Reference](../rewards-api.md) -- full API reference for decorators and `evaluate_rubric`
+- [Local Rollout Overview](./overview.md) -- repository structure and setup
diff --git a/docs/rollout/api-reference.md b/docs/remote-rollout/agent-loop.md
similarity index 98%
rename from docs/rollout/api-reference.md
rename to docs/remote-rollout/agent-loop.md
index 48732236..6ae20ea0 100644
--- a/docs/rollout/api-reference.md
+++ b/docs/remote-rollout/agent-loop.md
@@ -1,4 +1,4 @@
-# API Reference
+# Agent Loop Guide
Complete API documentation for the Osmosis Remote Rollout SDK.
@@ -714,3 +714,9 @@ MessageDict = Dict[str, Any]
SamplingParamsDict = Dict[str, Any]
# Example: {"temperature": 0.7, "max_tokens": 512}
```
+
+## See Also
+
+- [Remote Rollout Overview](./overview.md) - Quick start guide
+- [Architecture](./architecture.md) - Protocol design
+- [Examples](./examples.md) - Working code examples
diff --git a/docs/rollout/architecture.md b/docs/remote-rollout/architecture.md
similarity index 97%
rename from docs/rollout/architecture.md
rename to docs/remote-rollout/architecture.md
index c0acade1..97c083b8 100644
--- a/docs/rollout/architecture.md
+++ b/docs/remote-rollout/architecture.md
@@ -225,3 +225,9 @@ On server shutdown:
- Cleanup task is cancelled
- All running rollout tasks are cancelled
- Resources are released
+
+## See Also
+
+- [Remote Rollout Overview](./overview.md) - Quick start guide
+- [Agent Loop Guide](./agent-loop.md) - Complete API documentation
+- [Examples](./examples.md) - Working code examples
diff --git a/docs/rollout/deployment.md b/docs/remote-rollout/deployment.md
similarity index 100%
rename from docs/rollout/deployment.md
rename to docs/remote-rollout/deployment.md
diff --git a/docs/rollout/examples.md b/docs/remote-rollout/examples.md
similarity index 98%
rename from docs/rollout/examples.md
rename to docs/remote-rollout/examples.md
index d13444a5..a25ec7a1 100644
--- a/docs/rollout/examples.md
+++ b/docs/remote-rollout/examples.md
@@ -684,8 +684,12 @@ agent = get_agent_loop("calculator")
app = create_app(agent)
```
+## Example Repository
+
+For a complete, runnable project with tools, rewards, and server setup, see: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
+
## See Also
-- [Testing](./testing.md) — unit tests and mock trainer
-- [Deployment](./deployment.md) — Docker, health checks, production config
-- [API Reference](./api-reference.md) — endpoints, schemas, types
+- [Testing](./testing.md) -- unit tests and mock trainer
+- [Deployment](./deployment.md) -- Docker, health checks, production config
+- [Agent Loop Guide](./agent-loop.md) -- endpoints, schemas, types
diff --git a/docs/rollout/README.md b/docs/remote-rollout/overview.md
similarity index 92%
rename from docs/rollout/README.md
rename to docs/remote-rollout/overview.md
index a78d6878..81f437f4 100644
--- a/docs/rollout/README.md
+++ b/docs/remote-rollout/overview.md
@@ -11,6 +11,8 @@ The Remote Rollout SDK enables you to:
- Execute tools and collect training data (trajectories, token IDs, logprobs)
- Run agents as scalable HTTP servers
+For an alternative approach that requires no server infrastructure, see [Local Rollout](../local-rollout/overview.md).
+
## Installation
```bash
@@ -222,12 +224,16 @@ export OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS=120
export OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS=200
```
+## Example Repository
+
+See the complete working example: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
+
## Next Steps
- [Architecture](./architecture.md) - Understand the system design
-- [API Reference](./api-reference.md) - Complete API documentation
+- [Agent Loop Guide](./agent-loop.md) - Complete API documentation
- [Examples](./examples.md) - Full working examples
-- [Test Mode](./test-mode.md) - Local testing with external LLM providers
- [Testing](./testing.md) - Unit tests and mock trainer
- [Deployment](./deployment.md) - Docker, health checks, production config
-- [Dataset Format](./dataset-format.md) - Supported formats and required columns
+- [Dataset Format](../datasets.md) - Supported formats and required columns
+- [Test Mode](../test-mode.md) - Local testing with external LLM providers
diff --git a/docs/rollout/testing.md b/docs/remote-rollout/testing.md
similarity index 98%
rename from docs/rollout/testing.md
rename to docs/remote-rollout/testing.md
index d580803d..7a5e0533 100644
--- a/docs/rollout/testing.md
+++ b/docs/remote-rollout/testing.md
@@ -149,5 +149,5 @@ async def test_agent_run():
## See Also
- [Examples](./examples.md) -- agent implementations and utilities
-- [Test Mode](./test-mode.md) -- testing with cloud LLMs
+- [Test Mode](../test-mode.md) -- testing with cloud LLMs
- [Deployment](./deployment.md) -- production deployment
diff --git a/docs/rewards-api.md b/docs/rewards-api.md
new file mode 100644
index 00000000..2506acf7
--- /dev/null
+++ b/docs/rewards-api.md
@@ -0,0 +1,214 @@
+# Rewards API Reference
+
+API reference for reward function and rubric evaluation decorators. These are used in both Local Rollout and Remote Rollout modes for scoring LLM outputs during training.
+
+## @osmosis_reward
+
+Decorator for deterministic reward functions. Validates the function signature at decoration time and ensures a `float` return value at call time.
+
+### Required Signature
+
+```python
+from osmosis_ai import osmosis_reward
+
+@osmosis_reward
+def your_function(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
+ return float_score
+```
+
+### Parameters
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `solution_str` | `str` | Yes | The solution string to evaluate |
+| `ground_truth` | `str` | Yes | The correct/expected answer |
+| `extra_info` | `dict` | No (default `None`) | Optional dictionary for additional configuration |
+| `**kwargs` | - | Yes | Required for platform compatibility |
+
+### Return Value
+
+Must return a `float`. The decorator raises `TypeError` if the return type check fails.
+
+### Validation Rules
+
+- Function must accept exactly `solution_str`, `ground_truth`, `extra_info` as positional parameters
+- `extra_info` must have a default value of `None`
+- Must include `**kwargs`
+- Return type annotation must be `float`
+
+---
+
+## @osmosis_rubric
+
+Decorator for LLM-as-judge rubric functions. Validates the function signature and ensures required `extra_info` fields are present at call time.
+
+### Required Signature
+
+```python
+from osmosis_ai import osmosis_rubric
+
+@osmosis_rubric
+def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+ return float_score
+```
+
+### Parameters
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `solution_str` | `str` | Yes | The solution string to evaluate |
+| `ground_truth` | `str \| None` | Yes | The correct/expected answer (may be `None`) |
+| `extra_info` | `dict` | Yes (no default) | Configuration dict with provider, model, rubric |
+
+### Required `extra_info` Fields
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `provider` | `str` | Non-empty string identifying the judge provider |
+| `model` | `str` | Non-empty string naming the provider model |
+| `rubric` | `str` | Natural-language rubric instructions |
+| `api_key` or `api_key_env` | `str` | Raw API key or environment variable name |
+
+### Optional `extra_info` Fields
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `system_prompt` | `str` | Prepended to the provider's base system prompt |
+| `score_min` | `float` | Override minimum score bound |
+| `score_max` | `float` | Override maximum score bound |
+| `model_info_overrides` | `dict` | Merged into provider configuration |
+
+### Validation Rules
+
+- `extra_info` must be annotated as `dict` **without** a default value
+- `ground_truth` may be `None` (annotate as `Optional[str]` if needed)
+- Return type must be `float`
+
+---
+
+## evaluate_rubric()
+
+Evaluate a solution against a natural-language rubric using an external LLM judge via [LiteLLM](https://github.com/BerriAI/litellm).
+
+### Signature
+
+```python
+from osmosis_ai import evaluate_rubric
+
+score = evaluate_rubric(
+ rubric: str,
+ solution_str: str,
+ model_info: dict,
+ ground_truth: str = None,
+ metadata: dict = None,
+ return_details: bool = False,
+)
+```
+
+### Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `rubric` | `str` | required | Natural-language rubric instructions for the judge |
+| `solution_str` | `str` | required | The solution to evaluate |
+| `model_info` | `dict` | required | Provider configuration (see below) |
+| `ground_truth` | `str` | `None` | Optional reference answer |
+| `metadata` | `dict` | `None` | Optional structured context quoted in the judge prompt |
+| `return_details` | `bool` | `False` | Return full `RewardRubricRunResult` payload |
+
+### `model_info` Fields
+
+| Field | Type | Required | Description |
+|-------|------|----------|-------------|
+| `provider` | `str` | Yes | Provider name (e.g., `"openai"`, `"anthropic"`) |
+| `model` | `str` | Yes | Model name (e.g., `"gpt-5"`, `"claude-sonnet-4-5-20250929"`) |
+| `api_key` | `str` | No | Raw API key (mutually exclusive with `api_key_env`) |
+| `api_key_env` | `str` | No | Environment variable name for API key |
+| `score_min` | `float` | No | Minimum score bound (default `0.0`) |
+| `score_max` | `float` | No | Maximum score bound (default `1.0`) |
+| `system_prompt` | `str` | No | Optional context prepended to judge prompt |
+| `original_input` | `str` | No | Optional original user input for context |
+| `timeout` | `int` | No | Provider timeout in seconds |
+
+### Return Value
+
+- When `return_details=False`: Returns a `float` score clamped to `[score_min, score_max]`
+- When `return_details=True`: Returns `RewardRubricRunResult` with full provider response
+
+### Credential Resolution
+
+Credentials are resolved in this order:
+
+1. `api_key` in `model_info` (used directly if present and non-empty)
+2. `api_key_env` in `model_info` (environment variable name to look up)
+3. Provider default environment variable:
+
+| Provider | Default Environment Variable |
+|----------|------------------------------|
+| OpenAI | `OPENAI_API_KEY` |
+| Anthropic | `ANTHROPIC_API_KEY` |
+| Google Gemini | `GEMINI_API_KEY` |
+| xAI | `XAI_API_KEY` |
+| OpenRouter | `OPENROUTER_API_KEY` |
+| Cerebras | `CEREBRAS_API_KEY` |
+
+### Provider Architecture
+
+All provider routing is handled by [LiteLLM](https://github.com/BerriAI/litellm). Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
+
+Every provider returns a strict JSON object with `{"score": number, "explanation": string}`. The helper clamps the score into the configured range and validates the structure.
+
+---
+
+## Error Types
+
+### MissingAPIKeyError
+
+Raised when the required API key is not found in the environment.
+
+```python
+from osmosis_ai.rubric import MissingAPIKeyError
+
+try:
+ score = evaluate_rubric(...)
+except MissingAPIKeyError as e:
+ print(f"Missing key: {e}")
+ # Message explains which env var to export
+```
+
+### ProviderRequestError
+
+Raised when the LLM provider returns an error.
+
+```python
+from osmosis_ai.rubric import ProviderRequestError
+
+try:
+ score = evaluate_rubric(...)
+except ProviderRequestError as e:
+ print(f"Provider error: {e}")
+```
+
+### ModelNotFoundError
+
+Raised when the specified model identifier is not recognized by the provider.
+
+```python
+from osmosis_ai.rubric import ModelNotFoundError
+
+try:
+ score = evaluate_rubric(...)
+except ModelNotFoundError as e:
+ print(f"Model not found: {e}")
+ # Check provider dashboard for latest model names
+```
+
+> Provider model snapshot names change frequently. Check each vendor's dashboard for the latest identifier if you encounter a "model not found" error.
+
+---
+
+## See Also
+
+- [Reward Functions Guide](./local-rollout/reward-functions.md) -- usage guide with examples
+- [Reward Rubrics Guide](./local-rollout/reward-rubrics.md) -- usage guide with provider examples
+- [CLI Reference](./cli.md) -- `osmosis preview` and `osmosis eval-rubric` commands
diff --git a/docs/rewards.md b/docs/rewards.md
deleted file mode 100644
index 624ff1ab..00000000
--- a/docs/rewards.md
+++ /dev/null
@@ -1,189 +0,0 @@
-# Rewards & Rubrics
-
-The Osmosis SDK provides two mechanisms for scoring LLM outputs during training: **reward functions** for deterministic, code-based scoring, and **rubric evaluations** for LLM-as-judge assessments powered by external providers. Both integrate with the Osmosis training platform to drive reinforcement learning workflows.
-
-## @osmosis_reward
-
-All functions decorated with `@osmosis_reward` must have exactly this signature:
-
-```python
-from osmosis_ai import osmosis_reward
-
-@osmosis_reward
-def your_function(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- # Your reward logic here
- return float_score
-```
-
-### Parameters
-
-- **`solution_str: str`** -- The solution string to evaluate (required).
-- **`ground_truth: str`** -- The correct/expected answer (required).
-- **`extra_info: dict = None`** -- Optional dictionary for additional configuration.
-
-### Return Value
-
-- **`-> float`** -- Must return a float value representing the reward score.
-
-The decorator will raise a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
-
-## @osmosis_rubric
-
-Rubric functions decorated with `@osmosis_rubric` must match this signature:
-
-```python
-from osmosis_ai import osmosis_rubric
-
-@osmosis_rubric
-def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
- # Your rubric logic here
- return float_score
-```
-
-> The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
-
-### Required `extra_info` fields
-
-- **`provider`** -- Non-empty string identifying the judge provider.
-- **`model`** -- Non-empty string naming the provider model to call.
-- **`rubric`** -- Natural-language rubric instructions for the judge model.
-- **`api_key` / `api_key_env`** -- Supply either the raw key or the environment variable name that exposes it.
-
-### Optional `extra_info` fields
-
-- **`system_prompt`** -- Optional string prepended to the provider's base system prompt when invoking the judge; include it inside `extra_info` rather than as a separate argument.
-- **`score_min` / `score_max`** -- Optional numeric overrides for the expected score range.
-- **`model_info_overrides`** -- Optional dict merged into the provider configuration passed to the judge.
-
-Additional keys are passthrough and can be used for custom configuration. If you need to extend the provider payload (for example adding `api_key_env`), add a dict under `model_info_overrides` and it will be merged with the required `provider`/`model` pair before invoking `evaluate_rubric`. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
-
-> Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
-
-> Tip: When delegating to `evaluate_rubric`, pass the raw `solution_str` directly and include any extra context inside the `metadata` payload.
-
-## evaluate_rubric
-
-`evaluate_rubric` talks to hosted LLM providers through [LiteLLM](https://github.com/BerriAI/litellm), a unified interface supporting 100+ providers (OpenAI, Anthropic, Google Gemini, xAI, OpenRouter, Cerebras, Azure, Bedrock, Vertex AI, and more). Every provider returns a strict JSON object with `{"score": number, "explanation": string}`. The helper clamps the score into your configured range, validates the structure, and exposes the raw payload when `return_details=True`.
-
-### Basic Usage
-
-```python
-from osmosis_ai import evaluate_rubric
-
-solution = "The capital of France is Paris."
-
-rubric_score = evaluate_rubric(
- rubric="Assistant must mention the verified capital city.",
- solution_str=solution,
- model_info={
- "provider": "openai",
- "model": "gpt-5",
- "api_key_env": "OPENAI_API_KEY",
- },
- ground_truth="Paris",
-)
-
-print(rubric_score) # -> 1.0 (full payload available via return_details=True)
-```
-
-### Credentials
-
-Credentials are resolved from environment variables by default:
-
-- `OPENAI_API_KEY` for OpenAI
-- `ANTHROPIC_API_KEY` for Anthropic
-- `GEMINI_API_KEY` for Google Gemini
-- `XAI_API_KEY` for xAI
-- `OPENROUTER_API_KEY` for OpenRouter
-- `CEREBRAS_API_KEY` for Cerebras
-
-Override the environment variable name with `model_info={"api_key_env": "CUSTOM_ENV_NAME"}` when needed, or supply an inline secret with `model_info={"api_key": "sk-..."}` for ephemeral credentials. Missing API keys raise a `MissingAPIKeyError` that explains how to export the secret before trying again.
-
-`api_key` and `api_key_env` are mutually exclusive ways to provide the same credential. When `api_key` is present and non-empty it is used directly, skipping any environment lookup. Otherwise the resolver falls back to `api_key_env` (or the provider default) and pulls the value from your local environment with `os.getenv`.
-
-### Configuration Options
-
-`model_info` accepts additional rubric-specific knobs:
-
-- **`score_min` / `score_max`** -- Change the default `[0.0, 1.0]` scoring bounds.
-- **`system_prompt` / `original_input`** -- Provide optional context strings that will be quoted in the judging prompt.
-- **`timeout`** -- Customise the provider timeout in seconds.
-
-Pass `metadata={...}` to `evaluate_rubric` when you need structured context quoted in the judge prompt, and set `return_details=True` to receive the full `RewardRubricRunResult` payload (including the provider's raw response).
-
-### Provider Architecture
-
-All provider routing is handled by [LiteLLM](https://github.com/BerriAI/litellm). To use any supported provider, pass its name and model in `model_info`:
-
-```python
-result = evaluate_rubric(
- rubric="...",
- solution_str="...",
- model_info={"provider": "anthropic", "model": "claude-sonnet-4-5-20250929"},
-)
-```
-
-Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
-
-### Error Handling
-
-Remote failures surface as `ProviderRequestError` instances, with `ModelNotFoundError` reserved for missing model identifiers so you can retry with a new snapshot.
-
-> Provider model snapshot names change frequently. Check each vendor's dashboard for the latest identifier if you encounter a "model not found" error.
-
-## Examples
-
-See the [`examples/`](../examples/) directory for complete examples:
-
-```python
-@osmosis_reward
-def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- """Case-insensitive string matching with partial credit."""
- match = solution_str.lower().strip() == ground_truth.lower().strip()
-
- if extra_info and 'partial_credit' in extra_info:
- if not match and extra_info['partial_credit']:
- len_diff = abs(len(solution_str) - len(ground_truth))
- if len_diff <= 2:
- return 0.5
-
- return 1.0 if match else 0.0
-
-@osmosis_reward
-def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
- """Numeric comparison with configurable tolerance."""
- try:
- solution_num = float(solution_str.strip())
- truth_num = float(ground_truth.strip())
-
- tolerance = extra_info.get('tolerance', 0.01) if extra_info else 0.01
- return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
- except ValueError:
- return 0.0
-```
-
-- `examples/reward_functions.py` keeps local reward helpers that showcase the decorator contract without external calls.
-- `examples/rubric_functions.py` demonstrates `evaluate_rubric` with OpenAI, Anthropic, Gemini, xAI, OpenRouter, and Cerebras via LiteLLM's unified interface.
-- `examples/rubric_configs.yaml` bundles two rubric definitions with provider configuration and scoring bounds.
-- `examples/sample_data.jsonl` contains two rubric-aligned solution strings 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", "original_input": "...", "solution_str": "..."}
-{"conversation_id": "ticket-047", "rubric_id": "policy_grounding", "original_input": "...", "solution_str": "..."}
-```
-
-## See Also
-
-- [CLI Reference](./cli.md) -- `osmosis preview` and `osmosis eval-rubric` commands
-- [Rollout Evaluation](./rollout/eval.md) -- evaluating agents against datasets
diff --git a/docs/rollout/test-mode.md b/docs/test-mode.md
similarity index 91%
rename from docs/rollout/test-mode.md
rename to docs/test-mode.md
index 31fb3009..e0189c50 100644
--- a/docs/rollout/test-mode.md
+++ b/docs/test-mode.md
@@ -1,6 +1,6 @@
# Test Mode
-Test your agent implementations locally without TrainGate using external LLM providers via [LiteLLM](https://docs.litellm.ai/docs/providers).
+Test your agent implementations locally without TrainGate using external LLM providers via [LiteLLM](https://docs.litellm.ai/docs/providers). Works with both **Local Rollout** (MCP tools) and **Remote Rollout** (RolloutAgentLoop) agents.
> Breaking change: the legacy import path `osmosis_ai.rollout.test_mode` was removed.
> Use `osmosis_ai.rollout.eval.common` and `osmosis_ai.rollout.eval.test_mode` instead.
@@ -10,7 +10,7 @@ Test your agent implementations locally without TrainGate using external LLM pro
Test mode enables you to:
- Validate agent logic before deploying to training infrastructure
-- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp` (for git-sync users)
+- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp`
- Use 100+ LLM providers (OpenAI, Anthropic, Groq, Ollama, etc.)
- Run batch tests against datasets with detailed metrics
- Debug step-by-step with interactive mode
@@ -18,23 +18,25 @@ Test mode enables you to:
## Quick Start
-### CLI Usage
+### Testing with Remote Rollout Agent
+
+If you implement a `RolloutAgentLoop`, use `-m` to point at your agent module:
```bash
# Basic batch test with OpenAI
-osmosis test --agent my_agent:MyAgentLoop --dataset data.jsonl --model gpt-5-mini
+osmosis test -m my_agent:agent_loop -d data.jsonl --model gpt-5-mini
# Use Anthropic Claude
-osmosis test --agent my_agent:MyAgentLoop --dataset data.jsonl --model anthropic/claude-sonnet-4-5
+osmosis test -m my_agent:agent_loop -d data.jsonl --model anthropic/claude-sonnet-4-5
# Interactive debugging
-osmosis test --agent my_agent:MyAgentLoop --dataset data.jsonl --interactive
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
# Start at specific row
-osmosis test --agent my_agent:MyAgentLoop --dataset data.jsonl --interactive --row 5
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive --row 5
```
-### Git-Sync Users (MCP Tools)
+### Testing with Local Rollout MCP Tools
If you use the **git-sync** workflow and provide MCP tools (`@mcp.tool()`) instead of a `RolloutAgentLoop`, use `--mcp` to point at your MCP directory:
@@ -105,7 +107,7 @@ print(f"Total tokens: {results.total_tokens}")
## Dataset Format
-See [Dataset Format](./dataset-format.md) for supported formats and required columns.
+See [Dataset Format](./datasets.md) for supported formats and required columns.
---
@@ -125,8 +127,8 @@ osmosis test [OPTIONS]
| Option | Description |
|--------|-------------|
-| `-m, --module, --agent MODULE` | Module path to agent loop (format: `module:attribute`). Use for remote-rollout users who implement `RolloutAgentLoop`. |
-| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for git-sync users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
+| `-m, --module, --agent MODULE` | Module path to agent loop (format: `module:attribute`). Use for Remote Rollout users who implement `RolloutAgentLoop`. |
+| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for Local Rollout (git-sync) users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
### Model Options
@@ -173,13 +175,13 @@ See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for supported pr
### Examples
```bash
-# Test with GPT-5-mini (default)
+# Remote Rollout: test with GPT-5-mini (default)
osmosis test -m my_agent:agent_loop -d data.jsonl
-# Git-sync: test MCP tools without writing an AgentLoop
+# Local Rollout: test MCP tools
osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-# Git-sync: interactive debugging with MCP tools
+# Local Rollout: interactive debugging with MCP tools
osmosis test --mcp ./mcp -d data.jsonl --interactive
# Test with Claude
@@ -588,8 +590,8 @@ See [LiteLLM Environment Variables](https://docs.litellm.ai/docs/providers) for
## See Also
-- [Eval Mode](./eval.md) - Evaluate agents with eval functions and pass@k
-- [Architecture](./architecture.md) - System design overview
-- [API Reference](./api-reference.md) - Complete SDK API documentation
-- [Examples](./examples.md) - Working code examples
+- [Eval Mode](./eval-mode.md) - Evaluate agents with eval functions and pass@k
+- [Dataset Format](./datasets.md) - Supported formats and required columns
+- [Remote Rollout Examples](./remote-rollout/examples.md) - Working code examples
+- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) - MCP tool definition
- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
diff --git a/examples/README.md b/examples/README.md
index 3534314a..06fd2a98 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -2,7 +2,16 @@
This directory contains example usage of the osmosis-ai library.
-## Reward Functions (`reward_functions.py`)
+## Full Example Repos
+
+Complete, runnable example projects:
+
+- **Local Rollout**: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) — reward functions, rubrics, and MCP tools
+- **Remote Rollout**: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) — custom `RolloutAgentLoop` server
+
+## SDK API Examples
+
+### Reward Functions (`reward_functions.py`)
Demonstrates how to use the `@osmosis_reward` decorator with various reward function patterns:
@@ -41,7 +50,7 @@ PYTHONPATH=.. python reward_functions.py
This will run test cases through all the example reward functions and show their outputs.
-## Remote Rubric Evaluation (`rubric_functions.py`)
+### Remote Rubric Evaluation (`rubric_functions.py`)
Shows how to call `osmosis_ai.evaluate_rubric` against different hosted judge providers using their official Python SDKs. The example:
@@ -52,7 +61,7 @@ Shows how to call `osmosis_ai.evaluate_rubric` against different hosted judge pr
The helper uses [LiteLLM](https://github.com/BerriAI/litellm) under the hood; each example call only needs to supply `provider` and `model` because LiteLLM routes to the correct API automatically. Any provider supported by LiteLLM can be used by passing its name in `model_info`.
-## Rubric Configs and Dataset
+### Rubric Configs and Dataset
Use `rubric_configs.yaml` for a pair of ready-to-run rubric configurations (OpenAI and Anthropic) and `sample_data.jsonl` for two matching solution strings that exercise those rubrics. They are designed to work with `osmosis preview --path examples/` and can be adapted when building your own evaluation suites.
From 507f0de6b4e9f50a2271ff3d477c9aa916634ae1 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 14:05:10 -0800
Subject: [PATCH 36/60] Add logo to README for improved visual appeal
---
.github/osmosis-logo-dark.svg | 24 ++++++++++++++++++++++++
.github/osmosis-logo-light.svg | 24 ++++++++++++++++++++++++
README.md | 8 ++++++++
3 files changed, 56 insertions(+)
create mode 100644 .github/osmosis-logo-dark.svg
create mode 100644 .github/osmosis-logo-light.svg
diff --git a/.github/osmosis-logo-dark.svg b/.github/osmosis-logo-dark.svg
new file mode 100644
index 00000000..b32ee03a
--- /dev/null
+++ b/.github/osmosis-logo-dark.svg
@@ -0,0 +1,24 @@
+
diff --git a/.github/osmosis-logo-light.svg b/.github/osmosis-logo-light.svg
new file mode 100644
index 00000000..2ef43017
--- /dev/null
+++ b/.github/osmosis-logo-light.svg
@@ -0,0 +1,24 @@
+
diff --git a/README.md b/README.md
index 0b0677ac..b57b9d34 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,11 @@
+
+
+
+
+
+
+
+
# osmosis-ai
Python SDK for Osmosis AI training workflows. Supports two training modes with shared tooling for testing and evaluation.
From c8ef3efef591b15b8e5d4287a6708ce6a4d6cb93 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 14:14:55 -0800
Subject: [PATCH 37/60] Update pyproject.toml and README for Python version
support and badges
- Added support for Python versions 3.10 to 3.13 in pyproject.toml.
- Enhanced README with additional badges for platform, PyPI version, Python compatibility, license, and documentation links for better visibility.
---
README.md | 8 ++++++++
pyproject.toml | 4 ++++
2 files changed, 12 insertions(+)
diff --git a/README.md b/README.md
index b57b9d34..5ba1c3fa 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,14 @@
+
+
+
+
+
+
+
+
# osmosis-ai
Python SDK for Osmosis AI training workflows. Supports two training modes with shared tooling for testing and evaluation.
diff --git a/pyproject.toml b/pyproject.toml
index 44d92478..f5e38ed5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,6 +13,10 @@ authors = [
license = {file = "LICENSE"}
classifiers = [
"Programming Language :: Python :: 3",
+ "Programming Language :: Python :: 3.10",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
+ "Programming Language :: Python :: 3.13",
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
]
From 7fe2fd466809e0fc7313615f3f5a6fe7bb72e926 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 14:57:30 -0800
Subject: [PATCH 38/60] Enhance documentation and CLI options for improved user
experience
- Expanded the README to include a detailed Quick Start guide for Local Rollout, outlining installation steps and example commands.
- Updated prerequisites section in the README to clarify requirements for using the SDK.
- Revised CLI documentation to specify host and port options for the `osmosis serve` command.
- Added notes regarding breaking changes and command splits in eval and test modes.
- Improved clarity in various documentation sections, including agent loop and testing utilities.
---
README.md | 32 ++
SECURITY.md | 4 +-
docs/README.md | 6 +
docs/cli.md | 42 ++-
docs/configuration.md | 233 ++++++++++++++
docs/eval-mode.md | 32 +-
docs/local-rollout/overview.md | 2 +-
docs/local-rollout/reward-rubrics.md | 4 +-
docs/remote-rollout/agent-loop.md | 86 ++++-
docs/remote-rollout/architecture.md | 14 +-
docs/remote-rollout/deployment.md | 1 +
docs/remote-rollout/overview.md | 35 +-
docs/remote-rollout/testing.md | 51 +++
docs/rewards-api.md | 24 +-
docs/test-mode.md | 26 +-
docs/troubleshooting.md | 460 +++++++++++++++++++++++++++
examples/README.md | 4 +-
osmosis_ai/rollout/core/base.py | 9 +-
18 files changed, 977 insertions(+), 88 deletions(-)
create mode 100644 docs/configuration.md
create mode 100644 docs/troubleshooting.md
diff --git a/README.md b/README.md
index 5ba1c3fa..ce3db82e 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,32 @@
Python SDK for Osmosis AI training workflows. Supports two training modes with shared tooling for testing and evaluation.
+Osmosis AI is a platform for training LLMs with reinforcement learning. You define custom reward functions, LLM-as-judge rubrics, and agent tools -- then Osmosis handles the training loop on managed GPU clusters. This SDK provides everything you need to build and test those components locally, from `@osmosis_reward` decorators and MCP tool definitions to a full CLI for running agents against datasets before submitting training runs.
+
+## Quick Start
+
+The fastest way to get started is with **Local Rollout** using MCP tools. Clone the example repo and test it locally in under a minute:
+
+```bash
+# 1. Install the SDK with MCP support
+pip install "osmosis-ai[mcp]"
+
+# 2. Clone the example repo
+git clone https://github.com/Osmosis-AI/osmosis-git-sync-example.git
+cd osmosis-git-sync-example
+
+# 3. Set your LLM API key
+export OPENAI_API_KEY="sk-..."
+
+# 4. Run the agent against the sample dataset
+osmosis test --mcp ./mcp -d test_data.jsonl --model openai/gpt-4o-mini
+
+# 5. Try interactive mode to step through each LLM call
+osmosis test --mcp ./mcp -d test_data.jsonl --interactive
+```
+
+For detailed setup, see the [Local Rollout docs](docs/local-rollout/overview.md). For custom agent architectures, see [Remote Rollout](docs/remote-rollout/overview.md).
+
## Two Training Modes
Osmosis supports **Local Rollout** and **Remote Rollout** as parallel approaches to training with reinforcement learning:
@@ -32,6 +58,12 @@ Osmosis supports **Local Rollout** and **Remote Rollout** as parallel approaches
Requires Python 3.10 or newer. For development setup, see [CONTRIBUTING.md](CONTRIBUTING.md).
+### Prerequisites
+
+- **Python 3.10+**
+- **An LLM API key** (e.g., OpenAI, Anthropic, Groq) -- required for `osmosis test` and `osmosis eval`. See [supported providers](https://docs.litellm.ai/docs/providers).
+- **Osmosis account** (optional) -- needed for platform features like `osmosis login`, workspace management, and submitting training runs. Sign up at [osmosis.ai](https://osmosis.ai).
+
### pip
```bash
diff --git a/SECURITY.md b/SECURITY.md
index 7f43c8a1..54395e99 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -4,8 +4,8 @@
| Version | Supported |
|---------|--------------------|
-| Latest | :white_check_mark: |
-| < Latest | :x: |
+| Latest | Yes |
+| < Latest | No |
Only the latest published version of `osmosis-ai` receives security updates.
diff --git a/docs/README.md b/docs/README.md
index 8a4cc1a6..97d9f668 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -39,6 +39,12 @@ Used by both Local Rollout and Remote Rollout modes.
- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric` API details
- [CLI Reference](./cli.md) -- all `osmosis` commands
+- [Configuration](./configuration.md) -- environment variables, settings classes, and programmatic configuration
+- [Troubleshooting](./troubleshooting.md) -- common errors, debug tips, and resolutions
+
+## Contributing
+
+We welcome contributions! See the [Contributing Guide](../CONTRIBUTING.md) for development setup, coding standards, and pull request guidelines.
## Other
diff --git a/docs/cli.md b/docs/cli.md
index f313c350..e0158402 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -117,8 +117,8 @@ Start a RolloutServer for an agent loop implementation. The module path format i
osmosis login
osmosis serve -m my_agent:agent_loop
-# Specify port
-osmosis serve -m my_agent:agent_loop -p 8080
+# Specify host and port
+osmosis serve -m my_agent:agent_loop -H 127.0.0.1 -p 8080
# Local / container mode: skip Platform registration (no login required).
# NOTE: API key auth is still enabled by default.
@@ -130,20 +130,56 @@ osmosis serve -m my_agent:agent_loop --local
# Provide a stable API key (otherwise one is generated and printed on startup)
osmosis serve -m my_agent:agent_loop --skip-register --api-key "$MY_API_KEY"
+# Enable auto-reload for development
+osmosis serve -m my_agent:agent_loop --local --reload
+
+# Set log level and write execution traces to a directory
+osmosis serve -m my_agent:agent_loop --log-level debug --log ./traces
+
# Skip validation (not recommended)
osmosis serve -m my_agent:agent_loop --no-validate
```
-Note: The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This key is **not** the same as your `osmosis login` token (which is for authenticating with the Osmosis Platform), nor is it used for callbacks *from* your server back to TrainGate.
+**Options:**
+
+| Option | Description |
+|--------|-------------|
+| `-m`/`--module` | Module path to the agent loop (required) |
+| `-p`/`--port` | Port to bind to (default: 9000) |
+| `-H`/`--host` | Host to bind to (default: 0.0.0.0) |
+| `--skip-register` | Skip registering with Osmosis Platform (for local testing) |
+| `--local` | Local debug mode: disable API key auth and skip Platform registration |
+| `--api-key` | API key used by TrainGate to authenticate requests to this server |
+| `--no-validate` | Skip agent loop validation before starting |
+| `--reload` | Enable auto-reload for development |
+| `--log-level` | Uvicorn log level: debug, info, warning, error, critical (default: info) |
+| `--log DIR` | Enable logging and write per-rollout execution traces (as `{rollout_id}.jsonl`) to DIR |
+
+> **Note:** The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This key is **not** the same as your `osmosis login` token (which is for authenticating with the Osmosis Platform), nor is it used for callbacks *from* your server back to TrainGate.
+
+See [Remote Rollout Overview](./remote-rollout/overview.md) for architecture details and the full agent lifecycle.
### osmosis validate
Validate an agent loop before starting the server (checks tools, async run method, etc.):
```bash
+# Basic validation
osmosis validate -m my_agent:agent_loop
+
+# Verbose output with detailed warnings
+osmosis validate -m my_agent:agent_loop -v
```
+**Options:**
+
+| Option | Description |
+|--------|-------------|
+| `-m`/`--module` | Module path to the agent loop (required) |
+| `-v`/`--verbose` | Show detailed validation output including warnings |
+
+See [Remote Rollout Overview](./remote-rollout/overview.md) for architecture details and the full agent lifecycle.
+
## Rubric Tools
### osmosis preview
diff --git a/docs/configuration.md b/docs/configuration.md
new file mode 100644
index 00000000..7a35d082
--- /dev/null
+++ b/docs/configuration.md
@@ -0,0 +1,233 @@
+# Configuration
+
+The Osmosis rollout SDK is configured through environment variables, `.env` files, or programmatic Python calls. All settings use [Pydantic Settings](https://docs.pydantic.dev/latest/concepts/pydantic_settings/) when the `pydantic-settings` package is installed (included in the `server` extra).
+
+## Environment Variables
+
+Settings are loaded automatically from environment variables and from a `.env` file in the current working directory. The tables below list every recognized variable.
+
+### Client Settings (`OSMOSIS_ROLLOUT_CLIENT_*`)
+
+HTTP client configuration for communication between the RolloutServer and TrainGate.
+
+| Variable | Type | Default | Range | Description |
+|----------|------|---------|-------|-------------|
+| `OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS` | `float` | `300.0` | 1.0 -- 3600.0 | HTTP request timeout in seconds. |
+| `OSMOSIS_ROLLOUT_CLIENT_MAX_RETRIES` | `int` | `3` | 0 -- 10 | Maximum retry attempts for 5xx errors. |
+| `OSMOSIS_ROLLOUT_CLIENT_COMPLETE_ROLLOUT_RETRIES` | `int` | `2` | 0 -- 10 | Maximum retries for the completion callback (`/v1/rollout/completed`). |
+| `OSMOSIS_ROLLOUT_CLIENT_RETRY_BASE_DELAY` | `float` | `1.0` | 0.1 -- 60.0 | Base delay for exponential backoff in seconds. |
+| `OSMOSIS_ROLLOUT_CLIENT_RETRY_MAX_DELAY` | `float` | `30.0` | 1.0 -- 300.0 | Maximum delay between retries in seconds. |
+| `OSMOSIS_ROLLOUT_CLIENT_MAX_CONNECTIONS` | `int` | `100` | 1 -- 1000 | Maximum number of HTTP connections. |
+| `OSMOSIS_ROLLOUT_CLIENT_MAX_KEEPALIVE_CONNECTIONS` | `int` | `20` | 1 -- 100 | Maximum keepalive connections in the pool. |
+
+### Server Settings (`OSMOSIS_ROLLOUT_SERVER_*`)
+
+Server-side configuration for the RolloutServer process.
+
+| Variable | Type | Default | Range | Description |
+|----------|------|---------|-------|-------------|
+| `OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS` | `int` | `100` | 1 -- 10000 | Maximum number of concurrent rollouts. |
+| `OSMOSIS_ROLLOUT_SERVER_RECORD_TTL_SECONDS` | `float` | `3600.0` | 60.0 -- 86400.0 | How long to keep completed rollout records (seconds). |
+| `OSMOSIS_ROLLOUT_SERVER_CLEANUP_INTERVAL_SECONDS` | `float` | `60.0` | 10.0 -- 3600.0 | Interval for the cleanup task (seconds). |
+| `OSMOSIS_ROLLOUT_SERVER_REQUEST_TIMEOUT_SECONDS` | `float` | `600.0` | 10.0 -- 3600.0 | Timeout for individual requests (seconds). |
+| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_READINESS_TIMEOUT_SECONDS` | `float` | `10.0` | 1.0 -- 60.0 | Maximum wait for the server to become ready before platform registration. The server polls its own health endpoint to confirm readiness. |
+| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_READINESS_POLL_INTERVAL_SECONDS` | `float` | `0.2` | 0.05 -- 5.0 | Interval between health check polls during server readiness check. |
+| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_SHUTDOWN_TIMEOUT_SECONDS` | `float` | `30.0` | 1.0 -- 300.0 | Timeout for waiting for platform registration to complete during shutdown. |
+
+### Global Settings (`OSMOSIS_ROLLOUT_*`)
+
+Top-level settings that apply across the SDK.
+
+| Variable | Type | Default | Range | Description |
+|----------|------|---------|-------|-------------|
+| `OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES` | `int` | `1048576` (1 MB) | 1024 -- 104857600 (100 MB) | Maximum allowed size for rollout metadata in bytes. |
+
+## Settings Classes
+
+Configuration is organized into three Pydantic Settings classes. When `pydantic-settings` is installed, these classes automatically read from environment variables and `.env` files. Without `pydantic-settings`, they fall back to plain `pydantic.BaseModel` with only programmatic configuration.
+
+### `RolloutClientSettings`
+
+HTTP client settings for communication with TrainGate.
+
+```python
+from osmosis_ai.rollout.config import RolloutClientSettings
+
+client = RolloutClientSettings()
+print(client.timeout_seconds) # 300.0
+print(client.max_retries) # 3
+print(client.retry_base_delay) # 1.0
+```
+
+Env prefix: `OSMOSIS_ROLLOUT_CLIENT_`
+
+### `RolloutServerSettings`
+
+Server-side settings for the RolloutServer process.
+
+```python
+from osmosis_ai.rollout.config import RolloutServerSettings
+
+server = RolloutServerSettings()
+print(server.max_concurrent_rollouts) # 100
+print(server.request_timeout_seconds) # 600.0
+```
+
+Env prefix: `OSMOSIS_ROLLOUT_SERVER_`
+
+### `RolloutSettings`
+
+Top-level settings that aggregates `RolloutClientSettings` and `RolloutServerSettings` as nested objects.
+
+```python
+from osmosis_ai.rollout.config import RolloutSettings
+
+settings = RolloutSettings()
+print(settings.client.timeout_seconds) # 300.0
+print(settings.server.max_concurrent_rollouts) # 100
+print(settings.max_metadata_size_bytes) # 1048576
+```
+
+Env prefix: `OSMOSIS_ROLLOUT_`
+
+## Programmatic Configuration
+
+### `get_settings()`
+
+Returns the global settings singleton. On first call, settings are loaded from
+environment variables and `.env` files.
+
+```python
+from osmosis_ai.rollout.config import get_settings
+
+settings = get_settings()
+timeout = settings.client.timeout_seconds
+```
+
+### `configure()`
+
+Replaces the global settings singleton with a custom instance. Call this early
+in your application to override environment variable defaults.
+
+```python
+from osmosis_ai.rollout.config import (
+ configure,
+ RolloutSettings,
+ RolloutClientSettings,
+ RolloutServerSettings,
+)
+
+configure(RolloutSettings(
+ client=RolloutClientSettings(
+ timeout_seconds=120.0,
+ max_retries=5,
+ ),
+ server=RolloutServerSettings(
+ max_concurrent_rollouts=200,
+ ),
+))
+```
+
+### `reset_settings()`
+
+Resets the global settings singleton to `None`. Primarily used in tests to
+ensure a clean state.
+
+```python
+from osmosis_ai.rollout.config import reset_settings
+
+reset_settings()
+```
+
+## Metadata Size Limits
+
+Rollout requests include a `metadata` dict that is validated against a
+configurable size limit. The default limit is **1 MB**. The size is measured as
+the byte length of the JSON-serialized metadata.
+
+### `get_max_metadata_size_bytes()`
+
+Returns the current maximum metadata size limit in bytes. Thread-safe.
+
+```python
+from osmosis_ai.rollout.core.schemas import get_max_metadata_size_bytes
+
+limit = get_max_metadata_size_bytes() # 1048576 (1 MB)
+```
+
+### `set_max_metadata_size_bytes()`
+
+Sets the maximum metadata size limit in bytes. Thread-safe. The value must be
+positive; otherwise a `ValueError` is raised.
+
+```python
+from osmosis_ai.rollout.core.schemas import set_max_metadata_size_bytes
+
+# Increase to 2 MB
+set_max_metadata_size_bytes(2 * 1024 * 1024)
+```
+
+The limit is also configurable via the `max_metadata_size_bytes` field on
+`RolloutSettings` (env: `OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES`), though
+the runtime functions above provide thread-safe access used during request
+validation.
+
+## Server Constants
+
+The `osmosis_ai.rollout.server.serve` module defines the following defaults
+for `serve_agent_loop()`:
+
+| Constant | Value | Description |
+|----------|-------|-------------|
+| `DEFAULT_PORT` | `9000` | Default port the RolloutServer binds to. |
+| `DEFAULT_HOST` | `"0.0.0.0"` | Default host the RolloutServer binds to (all interfaces). |
+
+Override these when starting the server:
+
+```bash
+osmosis serve -m my_agent:agent_loop -p 8080
+```
+
+Or programmatically:
+
+```python
+from osmosis_ai.rollout.server import serve_agent_loop
+
+serve_agent_loop(my_agent, host="127.0.0.1", port=8080)
+```
+
+## Credential File Location
+
+Authentication tokens from `osmosis login` are stored at:
+
+```
+~/.config/osmosis/credentials.json
+```
+
+The configuration directory (`~/.config/osmosis/`) is created with mode `0700`
+(owner-only access). The credentials file is written with mode `0600`
+(owner read/write only).
+
+The credentials file uses a multi-workspace format. It stores credentials for
+each workspace you have logged in to, along with the currently active workspace
+name. Use `osmosis workspace list` to see all stored workspaces and
+`osmosis workspace switch ` to change the active one.
+
+## `.env` File Support
+
+Both the CLI and the settings classes load environment variables from a `.env`
+file in the current working directory via `python-dotenv`. This is useful for
+local development:
+
+```bash
+# .env
+OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS=120
+OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS=50
+OPENAI_API_KEY=sk-...
+```
+
+## See Also
+
+- [CLI Reference](./cli.md) -- all `osmosis` commands and options
+- [Troubleshooting](./troubleshooting.md) -- common errors and resolutions
+- [Deployment](./remote-rollout/deployment.md) -- production configuration for remote rollout
diff --git a/docs/eval-mode.md b/docs/eval-mode.md
index 101f0b4c..faa00e4d 100644
--- a/docs/eval-mode.md
+++ b/docs/eval-mode.md
@@ -17,7 +17,7 @@ Key capabilities:
- **Concurrent execution** with `--batch-size` for faster benchmarks
- Use LiteLLM providers (e.g., `openai/gpt-5-mini`) as the primary model
-> Command split (development-stage breaking change):
+> **Note:** Command split (development-stage breaking change):
> `osmosis eval-rubric` is for hosted rubric evaluation, while `osmosis eval` is for agent eval mode documented here.
## Quick Start
@@ -62,25 +62,7 @@ osmosis eval --mcp ./mcp -d data.jsonl \
--n 5 --batch-size 5
```
-The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools:
-
-```python
-# mcp/main.py
-from fastmcp import FastMCP
-
-mcp = FastMCP("my_tools")
-
-@mcp.tool()
-def add(a: float, b: float) -> str:
- """Add two numbers."""
- return str(a + b)
-
-@mcp.tool()
-def lookup(key: str) -> str:
- """Look up a value by key."""
- data = {"pi": "3.14159", "e": "2.71828"}
- return data.get(key, "not found")
-```
+The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) for the full folder structure and examples.
> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use one or the other.
@@ -718,8 +700,8 @@ See [Test Mode - Environment Variables](./test-mode.md#environment-variables) fo
## See Also
-- [Test Mode](./test-mode.md) - Test agent logic with external LLMs
-- [Dataset Format](./datasets.md) - Supported formats and required columns
-- [Remote Rollout Examples](./remote-rollout/examples.md) - Working code examples
-- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) - MCP tool definition
-- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
+- [Test Mode](./test-mode.md) -- Test agent logic with external LLMs
+- [Dataset Format](./datasets.md) -- Supported formats and required columns
+- [Remote Rollout Examples](./remote-rollout/examples.md) -- Working code examples
+- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) -- MCP tool definition
+- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) -- Supported LLM providers
diff --git a/docs/local-rollout/overview.md b/docs/local-rollout/overview.md
index c5c05f0b..b66ee5a4 100644
--- a/docs/local-rollout/overview.md
+++ b/docs/local-rollout/overview.md
@@ -1,6 +1,6 @@
# Local Rollout
-Local Rollout is one of two training modes supported by the Osmosis platform. In this mode, Osmosis manages the entire agent loop -- you provide **reward functions**, **rubric evaluators**, and optionally **MCP tools** via a GitHub-synced repository. The training infrastructure handles LLM inference, tool execution, and trajectory collection automatically.
+Local Rollout is one of two training modes supported by the Osmosis platform. In this mode, Osmosis manages the entire agent loop -- you provide **reward functions**, **rubric evaluators**, and optionally **MCP tools** via a git-sync repository. The training infrastructure handles LLM inference, tool execution, and trajectory collection automatically.
## When to Choose Local Rollout
diff --git a/docs/local-rollout/reward-rubrics.md b/docs/local-rollout/reward-rubrics.md
index 83b85054..ca049114 100644
--- a/docs/local-rollout/reward-rubrics.md
+++ b/docs/local-rollout/reward-rubrics.md
@@ -15,7 +15,7 @@ def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -
return float_score
```
-> The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
+> **Note:** The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
### Required `extra_info` Fields
@@ -32,7 +32,7 @@ def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -
Additional keys are passthrough and can be used for custom configuration. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
-> Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
+> **Note:** Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
## evaluate_rubric
diff --git a/docs/remote-rollout/agent-loop.md b/docs/remote-rollout/agent-loop.md
index 6ae20ea0..f5b42742 100644
--- a/docs/remote-rollout/agent-loop.md
+++ b/docs/remote-rollout/agent-loop.md
@@ -27,9 +27,26 @@ class MyAgent(RolloutAgentLoop):
|-----------|------|-------------|
| `name` | `str` | Unique identifier for this agent loop (required) |
+**Subclass Validation (`__init_subclass__`):**
+
+When you define a concrete (non-abstract) subclass of `RolloutAgentLoop`, Python automatically validates at class definition time that the `name` class attribute is defined and non-empty. If it is missing or falsy, a `TypeError` is raised immediately:
+
+```python
+# This raises TypeError at class definition time:
+class BadAgent(RolloutAgentLoop):
+ # Missing 'name' attribute!
+ def get_tools(self, request):
+ return []
+ async def run(self, ctx):
+ return ctx.complete([])
+# TypeError: Agent loop class BadAgent must define a 'name' class attribute
+```
+
+Abstract subclasses (those with remaining abstract methods) are not validated, so you can create intermediate base classes without defining `name`.
+
**Abstract Methods:**
-#### `get_tools(request: RolloutRequest) -> List[OpenAIFunctionToolSchema]`
+#### `get_tools(request: RolloutRequest) -> list[OpenAIFunctionToolSchema]`
Return tools available for this rollout.
@@ -47,6 +64,23 @@ Execute the agent loop.
- **Returns:** `RolloutResult` with final status and messages
- **Notes:** Messages must be append-only; never modify previous messages
+**Convenience Methods:**
+
+#### `get_default_tools() -> list[OpenAIFunctionToolSchema]`
+
+Return the default tool list for discovery and validation purposes, without requiring a real `RolloutRequest`.
+
+Internally, this calls `get_tools()` with a synthetic request (`rollout_id="discovery"`, empty messages and params). This is used by the validation framework (`validate_agent_loop()`) and platform registration to discover the agent's tools without an active rollout.
+
+```python
+agent = MyAgent()
+tools = agent.get_default_tools()
+print(f"Agent provides {len(tools)} tools")
+```
+
+- **Returns:** list of `OpenAIFunctionToolSchema` objects
+- **Notes:** If your `get_tools()` returns different tools based on request metadata, `get_default_tools()` returns the tools for an empty/default request.
+
---
### RolloutContext
@@ -57,8 +91,8 @@ Execution context provided to the agent loop.
@dataclass
class RolloutContext:
request: RolloutRequest
- tools: List[OpenAIFunctionToolSchema]
- llm: OsmosisLLMClient
+ tools: list[OpenAIFunctionToolSchema]
+ llm: LLMClientProtocol
```
**Attributes:**
@@ -66,8 +100,8 @@ class RolloutContext:
| Attribute | Type | Description |
|-----------|------|-------------|
| `request` | `RolloutRequest` | Original rollout request |
-| `tools` | `List[OpenAIFunctionToolSchema]` | Tools returned by `get_tools()` |
-| `llm` | `OsmosisLLMClient` | HTTP client for LLM calls |
+| `tools` | `list[OpenAIFunctionToolSchema]` | Tools returned by `get_tools()` |
+| `llm` | `LLMClientProtocol` | LLM client for chat completions. In production (served via `create_app()` / `serve_agent_loop()`), this is an `OsmosisLLMClient` that calls back to TrainGate. In test/eval mode (e.g., `osmosis test`), this may be an `ExternalLLMClient` that routes requests to an external provider (OpenAI, Anthropic, etc.) via LiteLLM. |
**Methods:**
@@ -79,14 +113,23 @@ Shorthand for `self.llm.chat_completions()`.
result = await ctx.chat(messages, temperature=0.7)
```
-#### `complete(final_messages, finish_reason="stop") -> RolloutResult`
+#### `complete(final_messages, finish_reason="stop", reward=None) -> RolloutResult`
Create a successful completion result.
```python
return ctx.complete(messages, finish_reason="stop")
+
+# With reward
+return ctx.complete(messages, finish_reason="stop", reward=1.0)
```
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `final_messages` | `List[Dict[str, Any]]` | required | Final conversation messages |
+| `finish_reason` | `str` | `"stop"` | Why the rollout ended |
+| `reward` | `float \| None` | `None` | Optional precomputed trajectory reward score |
+
#### `error(error_message, final_messages=None) -> RolloutResult`
Create an error result.
@@ -143,6 +186,7 @@ class RolloutResult:
finish_reason: str
error_message: Optional[str] = None
metrics: Optional[RolloutMetrics] = None
+ reward: Optional[float] = None # Precomputed trajectory reward
```
---
@@ -213,6 +257,7 @@ Notify TrainGate that rollout is complete.
| `finish_reason` | `str` | `"stop"` | Why rollout ended |
| `error_message` | `Optional[str]` | `None` | Error description |
| `metrics` | `Optional[RolloutMetrics]` | `None` | Execution metrics |
+| `reward` | `Optional[float]` | `None` | Precomputed trajectory reward score |
#### `get_metrics() -> RolloutMetrics`
@@ -310,6 +355,21 @@ app = create_app(
record_ttl_seconds=3600,
debug_dir="./rollout_logs", # Optional: enable debug logging
)
+
+# Full options
+app = create_app(
+ agent_loop,
+ max_concurrent=100,
+ record_ttl_seconds=3600,
+ settings=None,
+ credentials=my_credentials, # For platform registration
+ server_host="0.0.0.0",
+ server_port=9000,
+ api_key="my-secret-key",
+ debug_dir="./rollout_logs",
+ on_startup=my_startup_callback,
+ on_shutdown=my_shutdown_callback,
+)
```
**Parameters:**
@@ -317,10 +377,16 @@ app = create_app(
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `agent_loop` | `RolloutAgentLoop` | required | Your agent implementation |
-| `max_concurrent` | `int` | `100` | Max concurrent rollouts |
-| `record_ttl_seconds` | `float` | `3600` | TTL for completed records |
+| `max_concurrent` | `int \| None` | `None` | Max concurrent rollouts (defaults to settings) |
+| `record_ttl_seconds` | `float \| None` | `None` | TTL for completed records (defaults to settings) |
| `settings` | `Optional[RolloutSettings]` | `None` | Global settings override (server/logging/tracing) |
+| `credentials` | `Optional[WorkspaceCredentials]` | `None` | Workspace credentials for platform registration. If `None`, registration is skipped |
+| `server_host` | `Optional[str]` | `None` | Host the server is bound to (used for platform registration) |
+| `server_port` | `Optional[int]` | `None` | Port the server is listening on (used for platform registration) |
+| `api_key` | `Optional[str]` | `None` | API key for authenticating incoming requests. If provided, requests must include `Authorization: Bearer ` |
| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging. If set, each rollout writes execution traces to `{debug_dir}/{rollout_id}.jsonl`. Note: when using `serve_agent_loop()` or CLI, a timestamped subdirectory is created automatically |
+| `on_startup` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async callback to run during application startup (e.g., warming caches, starting background services) |
+| `on_shutdown` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async callback to run during application shutdown (e.g., stopping services, releasing resources) |
**Returns:** `FastAPI` application
@@ -329,7 +395,8 @@ app = create_app(
| Endpoint | Method | Status | Description |
|----------|--------|--------|-------------|
| `/v1/rollout/init` | POST | 202 | Accept rollout request |
-| `/health` | GET | 200 | Health check |
+| `/health` | GET | 200 | Health check (unauthenticated) |
+| `/platform/health` | GET | 200 | Authenticated health check for Osmosis Platform. Requires `Authorization: Bearer `. Returns 404 if `api_key` is not configured (e.g., `local_debug` mode). Used by the platform to verify reachability and API key correctness. |
---
@@ -537,6 +604,7 @@ class RolloutResponse(BaseModel):
final_messages: List[MessageDict] = []
finish_reason: Optional[str] = None
error_message: Optional[str] = None
+ reward: Optional[float] = None # Precomputed trajectory reward
metrics: Optional[RolloutMetrics] = None
extra_fields: Dict[str, Any] = {}
```
diff --git a/docs/remote-rollout/architecture.md b/docs/remote-rollout/architecture.md
index 97c083b8..3d1a516e 100644
--- a/docs/remote-rollout/architecture.md
+++ b/docs/remote-rollout/architecture.md
@@ -103,7 +103,7 @@ class RolloutAgentLoop(ABC):
name: str # Unique identifier
@abstractmethod
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
+ def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
"""Return tools for this rollout."""
@abstractmethod
@@ -121,11 +121,11 @@ Execution context provided to the agent:
@dataclass
class RolloutContext:
request: RolloutRequest # Original request
- tools: List[OpenAIFunctionToolSchema] # Available tools
- llm: OsmosisLLMClient # LLM client
+ tools: list[OpenAIFunctionToolSchema] # Available tools
+ llm: LLMClientProtocol # LLM client (OsmosisLLMClient in production)
async def chat(self, messages, **kwargs) -> CompletionsResult
- def complete(self, messages, finish_reason="stop") -> RolloutResult
+ def complete(self, messages, finish_reason="stop", reward=None) -> RolloutResult
def error(self, message, final_messages=None) -> RolloutResult
def record_tool_call(self, latency_ms=0.0) -> None
```
@@ -228,6 +228,6 @@ On server shutdown:
## See Also
-- [Remote Rollout Overview](./overview.md) - Quick start guide
-- [Agent Loop Guide](./agent-loop.md) - Complete API documentation
-- [Examples](./examples.md) - Working code examples
+- [Remote Rollout Overview](./overview.md) -- Quick start guide
+- [Agent Loop Guide](./agent-loop.md) -- Complete API documentation
+- [Examples](./examples.md) -- Working code examples
diff --git a/docs/remote-rollout/deployment.md b/docs/remote-rollout/deployment.md
index 2eb3d528..8e0ab565 100644
--- a/docs/remote-rollout/deployment.md
+++ b/docs/remote-rollout/deployment.md
@@ -103,6 +103,7 @@ logging.getLogger("httpx").setLevel(logging.WARNING)
## See Also
+- [Configuration](../configuration.md) -- environment variables and SDK settings
- [Examples](./examples.md) -- agent implementations
- [Testing](./testing.md) -- unit tests and mock trainer
- [Architecture](./architecture.md) -- protocol design
diff --git a/docs/remote-rollout/overview.md b/docs/remote-rollout/overview.md
index 81f437f4..30d2971e 100644
--- a/docs/remote-rollout/overview.md
+++ b/docs/remote-rollout/overview.md
@@ -33,6 +33,8 @@ pip install osmosis-ai[full]
Create a class that inherits from `RolloutAgentLoop` and implements two methods:
```python
+import json
+
from osmosis_ai.rollout import (
RolloutAgentLoop,
RolloutContext,
@@ -57,6 +59,12 @@ class MyAgentLoop(RolloutAgentLoop):
)
]
+ async def execute_tool(self, name: str, args: dict) -> str:
+ """Execute a tool and return the result as a string."""
+ if name == "search":
+ return f"Search results for '{args.get('query', '')}'"
+ return f"Unknown tool: {name}"
+
async def run(self, ctx: RolloutContext) -> RolloutResult:
"""Execute the agent loop."""
messages = list(ctx.request.messages)
@@ -72,8 +80,17 @@ class MyAgentLoop(RolloutAgentLoop):
# Execute tools and add results
for tool_call in result.tool_calls:
- tool_result = await self.execute_tool(tool_call)
- messages.append(tool_result)
+ tool_name = tool_call["function"]["name"]
+ tool_args = json.loads(tool_call["function"]["arguments"])
+
+ # Execute tool logic (replace with your implementation)
+ tool_result = await self.execute_tool(tool_name, tool_args)
+
+ messages.append({
+ "role": "tool",
+ "content": tool_result,
+ "tool_call_id": tool_call["id"],
+ })
ctx.record_tool_call()
return ctx.complete(messages)
@@ -230,10 +247,10 @@ See the complete working example: [osmosis-remote-rollout-example](https://githu
## Next Steps
-- [Architecture](./architecture.md) - Understand the system design
-- [Agent Loop Guide](./agent-loop.md) - Complete API documentation
-- [Examples](./examples.md) - Full working examples
-- [Testing](./testing.md) - Unit tests and mock trainer
-- [Deployment](./deployment.md) - Docker, health checks, production config
-- [Dataset Format](../datasets.md) - Supported formats and required columns
-- [Test Mode](../test-mode.md) - Local testing with external LLM providers
+- [Architecture](./architecture.md) -- Understand the system design
+- [Agent Loop Guide](./agent-loop.md) -- Complete API documentation
+- [Examples](./examples.md) -- Full working examples
+- [Testing](./testing.md) -- Unit tests and mock trainer
+- [Deployment](./deployment.md) -- Docker, health checks, production config
+- [Dataset Format](../datasets.md) -- Supported formats and required columns
+- [Test Mode](../test-mode.md) -- Local testing with external LLM providers
diff --git a/docs/remote-rollout/testing.md b/docs/remote-rollout/testing.md
index 7a5e0533..b076efa2 100644
--- a/docs/remote-rollout/testing.md
+++ b/docs/remote-rollout/testing.md
@@ -146,6 +146,57 @@ async def test_agent_run():
assert result.status == "COMPLETED"
```
+## Testing Utilities Reference
+
+The `osmosis_ai.rollout.testing` module provides the following public utilities for testing agent loops without a real TrainGate server.
+
+### `create_mock_trainer_app(tracker=None, tool_call_generator=None)`
+
+Creates a mock trainer FastAPI application that implements the same HTTP endpoints as a real TrainGate server:
+
+- `POST /v1/chat/completions` -- returns deterministic LLM responses with fake token IDs
+- `POST /v1/rollout/completed` -- accepts completion callbacks
+- `GET /v1/rollout/completed/{rollout_id}` -- queries completed rollouts
+- `GET /health` -- health check
+
+By default, the mock generates tool calls when it detects calculator-related keywords (e.g. "add", "calculate", "multiply") in the user message. Pass a custom `tool_call_generator` function to override this behavior.
+
+### `RolloutCompletionTracker`
+
+Thread-safe tracker that captures `/v1/rollout/completed` callbacks during tests.
+
+| Attribute / Method | Description |
+|--------------------|-------------|
+| `event` | `threading.Event` that is set when a completion is received |
+| `responses` | List of captured completion response dicts |
+| `record(response)` | Record a response and signal the event |
+| `clear()` | Clear recorded responses and reset the event |
+| `wait(timeout=5.0)` | Block until a completion is received; returns `True` on success, `False` on timeout |
+
+### `patch_httpx_for_mock_trainer(client, monkeypatch)`
+
+Patches `httpx.AsyncClient.post` so that any requests to `/v1/chat/completions` or `/v1/rollout/completed` are routed to the mock trainer `TestClient` instead of making real HTTP calls. All other requests pass through unchanged.
+
+```python
+@pytest.fixture
+def mock_trainer(monkeypatch):
+ tracker = RolloutCompletionTracker()
+ app = create_mock_trainer_app(tracker=tracker)
+ client = TestClient(app)
+ patch_httpx_for_mock_trainer(client, monkeypatch)
+ return client, tracker
+```
+
+### `fake_token_ids(text)`
+
+Generates deterministic fake token IDs for testing. Returns a list of sequential integers, one per character in the input text (e.g. `fake_token_ids("hello")` returns `[0, 1, 2, 3, 4]`).
+
+### `fake_prompt_token_ids(messages)`
+
+Generates deterministic fake prompt token IDs for testing. The token count scales with the number of messages to simulate realistic prompt growth. Returns `list(range(10 * max(1, len(messages))))`.
+
+> **Note:** These fake token ID functions produce deterministic output suitable for snapshot testing but do not correspond to any real tokenizer.
+
## See Also
- [Examples](./examples.md) -- agent implementations and utilities
diff --git a/docs/rewards-api.md b/docs/rewards-api.md
index 2506acf7..cba1a507 100644
--- a/docs/rewards-api.md
+++ b/docs/rewards-api.md
@@ -33,8 +33,8 @@ Must return a `float`. The decorator raises `TypeError` if the return type check
- Function must accept exactly `solution_str`, `ground_truth`, `extra_info` as positional parameters
- `extra_info` must have a default value of `None`
-- Must include `**kwargs`
-- Return type annotation must be `float`
+- Must include `**kwargs` for platform compatibility
+- Return value must be a `float` (checked at call time)
---
@@ -128,12 +128,20 @@ score = evaluate_rubric(
| `score_max` | `float` | No | Maximum score bound (default `1.0`) |
| `system_prompt` | `str` | No | Optional context prepended to judge prompt |
| `original_input` | `str` | No | Optional original user input for context |
-| `timeout` | `int` | No | Provider timeout in seconds |
+| `timeout` | `float` | No | Provider timeout in seconds |
+| `reasoning_effort` | `str \| None` | No | Reasoning effort hint passed to the provider (e.g., `"low"`, `"medium"`, `"high"`). Silently dropped for models that do not support it. |
### Return Value
- When `return_details=False`: Returns a `float` score clamped to `[score_min, score_max]`
-- When `return_details=True`: Returns `RewardRubricRunResult` with full provider response
+- When `return_details=True`: Returns a `RewardRubricRunResult` dict with the following structure:
+
+```python
+class RewardRubricRunResult(TypedDict):
+ score: float # The clamped rubric score
+ explanation: str # The judge model's explanation for the score
+ raw: Any # The full raw response from the LLM provider
+```
### Credential Resolution
@@ -167,7 +175,7 @@ Every provider returns a strict JSON object with `{"score": number, "explanation
Raised when the required API key is not found in the environment.
```python
-from osmosis_ai.rubric import MissingAPIKeyError
+from osmosis_ai import MissingAPIKeyError
try:
score = evaluate_rubric(...)
@@ -181,7 +189,7 @@ except MissingAPIKeyError as e:
Raised when the LLM provider returns an error.
```python
-from osmosis_ai.rubric import ProviderRequestError
+from osmosis_ai import ProviderRequestError
try:
score = evaluate_rubric(...)
@@ -191,10 +199,10 @@ except ProviderRequestError as e:
### ModelNotFoundError
-Raised when the specified model identifier is not recognized by the provider.
+Raised when the specified model identifier is not recognized by the provider. Subclass of `ProviderRequestError`.
```python
-from osmosis_ai.rubric import ModelNotFoundError
+from osmosis_ai import ModelNotFoundError
try:
score = evaluate_rubric(...)
diff --git a/docs/test-mode.md b/docs/test-mode.md
index e0189c50..dcb2a138 100644
--- a/docs/test-mode.md
+++ b/docs/test-mode.md
@@ -2,7 +2,7 @@
Test your agent implementations locally without TrainGate using external LLM providers via [LiteLLM](https://docs.litellm.ai/docs/providers). Works with both **Local Rollout** (MCP tools) and **Remote Rollout** (RolloutAgentLoop) agents.
-> Breaking change: the legacy import path `osmosis_ai.rollout.test_mode` was removed.
+> **Note:** Breaking change: the legacy import path `osmosis_ai.rollout.test_mode` was removed.
> Use `osmosis_ai.rollout.eval.common` and `osmosis_ai.rollout.eval.test_mode` instead.
## Overview
@@ -55,19 +55,7 @@ osmosis test --mcp ./mcp -d data.jsonl --model gpt-5-mini \
--limit 5 --temperature 0.7 -o results.json
```
-The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools:
-
-```python
-# mcp/main.py
-from fastmcp import FastMCP
-
-mcp = FastMCP("my_tools")
-
-@mcp.tool()
-def add(a: float, b: float) -> str:
- """Add two numbers."""
- return str(a + b)
-```
+The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) for the full folder structure and examples.
> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use one or the other.
@@ -590,8 +578,8 @@ See [LiteLLM Environment Variables](https://docs.litellm.ai/docs/providers) for
## See Also
-- [Eval Mode](./eval-mode.md) - Evaluate agents with eval functions and pass@k
-- [Dataset Format](./datasets.md) - Supported formats and required columns
-- [Remote Rollout Examples](./remote-rollout/examples.md) - Working code examples
-- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) - MCP tool definition
-- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) - Supported LLM providers
+- [Eval Mode](./eval-mode.md) -- Evaluate agents with eval functions and pass@k
+- [Dataset Format](./datasets.md) -- Supported formats and required columns
+- [Remote Rollout Examples](./remote-rollout/examples.md) -- Working code examples
+- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) -- MCP tool definition
+- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) -- Supported LLM providers
diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md
new file mode 100644
index 00000000..88095d4a
--- /dev/null
+++ b/docs/troubleshooting.md
@@ -0,0 +1,460 @@
+# Troubleshooting
+
+Common errors and their resolutions when working with the Osmosis SDK.
+
+## Installation Issues
+
+### Missing extras
+
+The SDK ships optional dependency groups. If you see an `ImportError` for a
+package that should be available, you likely need to install the correct extra.
+
+| Error | Install command |
+|-------|----------------|
+| `No module named 'fastmcp'` | `pip install osmosis-ai[mcp]` |
+| `No module named 'fastapi'` / `No module named 'uvicorn'` | `pip install osmosis-ai[server]` |
+| `No module named 'pydantic_settings'` | `pip install osmosis-ai[server]` |
+| `No module named 'pyarrow'` | `pip install osmosis-ai[server]` (or `pip install pyarrow`) |
+| `No module named 'rich'` | `pip install osmosis-ai[server]` |
+| `No module named 'litellm'` | `pip install osmosis-ai` (included in core dependencies) |
+
+To install everything at once:
+
+```bash
+pip install osmosis-ai[full] # server + mcp
+pip install osmosis-ai[dev] # full + testing/linting tools
+```
+
+### Python version requirements
+
+The SDK requires **Python 3.10 or later** (3.10, 3.11, 3.12, 3.13). If you see
+syntax errors or compatibility issues, verify your Python version:
+
+```bash
+python --version
+```
+
+## Authentication Issues
+
+Credentials are saved to `~/.config/osmosis/credentials.json` after a
+successful `osmosis login`.
+
+### `osmosis login` fails
+
+1. **Browser does not open** -- pass `--no-browser` to print the authentication
+ URL so you can open it manually.
+2. **Network / firewall** -- the login flow requires outbound HTTPS to the
+ Osmosis platform. Ensure your network allows it.
+3. **Force re-login** -- if your session is in an inconsistent state, run:
+
+ ```bash
+ osmosis login --force
+ ```
+
+### Token expiration
+
+Tokens have an expiration time. When a token expires, commands that require
+authentication will fail silently or report that you are not logged in.
+
+```
+Not logged in. Please run 'osmosis login' first, or use skip_register=True for local testing.
+```
+
+Resolution:
+
+```bash
+osmosis login
+```
+
+To check your current session:
+
+```bash
+osmosis whoami
+```
+
+### Multiple workspaces
+
+If you are logged in to the wrong workspace, commands may fail with permission
+errors. List and switch workspaces with:
+
+```bash
+osmosis workspace list
+osmosis workspace switch
+```
+
+## Reward Function Errors
+
+### Signature validation errors from `@osmosis_reward`
+
+The `@osmosis_reward` decorator enforces the classic Osmosis signature when the
+first parameter is named `solution_str`:
+
+```
+(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float
+```
+
+Common `TypeError` messages and their causes:
+
+| Error message | Cause |
+|---------------|-------|
+| `Function must have at least 3 parameters, got ` | The function has fewer than three parameters. |
+| `First parameter 'solution_str' must be annotated as str, got ` | Missing or wrong type annotation on `solution_str`. |
+| `Second parameter must be named 'ground_truth', got ''` | The second parameter has the wrong name. |
+| `Second parameter 'ground_truth' must be annotated as str, got ` | Missing or wrong type annotation on `ground_truth`. |
+| `Third parameter must be named 'extra_info', got ''` | The third parameter has the wrong name. |
+| `Third parameter 'extra_info' must be annotated as dict or dict \| None, got ` | Wrong annotation on `extra_info`. Use `dict` or `dict \| None`. |
+| `Third parameter 'extra_info' must have a default value of None` | `extra_info` does not have a default value. |
+
+### Return type must be float
+
+When using the classic signature, the decorated function **must** return a plain
+`float`. Any other return type raises:
+
+```
+TypeError: Function must return a float, got
+```
+
+Ensure you return `float(value)` rather than `int`, `Decimal`, or another
+numeric type.
+
+## Rubric Evaluation Errors
+
+These errors originate from `osmosis_ai.rubric_eval` and
+`osmosis_ai.rubric_types`.
+
+### MissingAPIKeyError
+
+Raised when the LLM provider API key cannot be found. The error message
+includes a hint showing the expected environment variable.
+
+Example:
+
+```
+MissingAPIKeyError: Environment variable 'OPENAI_API_KEY' is not set.
+Export it with your openai API key before calling evaluate_rubric.
+Set the required API key before running:
+
+ export OPENAI_API_KEY="..."
+```
+
+Resolution -- set the appropriate environment variable for your provider:
+
+| Provider | Environment variable |
+|----------|---------------------|
+| openai | `OPENAI_API_KEY` |
+| anthropic | `ANTHROPIC_API_KEY` |
+| xai | `XAI_API_KEY` |
+| gemini / google | `GEMINI_API_KEY` |
+| openrouter | `OPENROUTER_API_KEY` |
+| cerebras | `CEREBRAS_API_KEY` |
+| azure | `AZURE_API_KEY` |
+| bedrock | `AWS_ACCESS_KEY_ID` |
+| vertex_ai | `GOOGLE_APPLICATION_CREDENTIALS` |
+
+You can also provide the key directly in `model_info` via the `api_key` or
+`api_key_env` fields.
+
+### ProviderRequestError
+
+Raised when the LLM provider call fails. The error includes the provider name,
+model name, and a detail string.
+
+```
+ProviderRequestError: Provider 'openai' request for model 'gpt-5-mini' failed.
+```
+
+Common causes:
+
+- **Authentication failure** -- your API key is invalid or expired.
+- **Rate limiting** -- you have exceeded the provider's rate limit. Wait and
+ retry.
+- **Timeout** -- the request took too long. Try increasing the `timeout`
+ parameter or use a faster model.
+- **Connection error** -- network issue between you and the provider. Check
+ connectivity.
+- **Invalid JSON response** -- the model returned content that could not be
+ parsed. Refine rubric instructions so the model returns valid JSON.
+
+### ModelNotFoundError
+
+A subclass of `ProviderRequestError` raised when the requested model does not
+exist.
+
+```
+ModelNotFoundError: Provider 'openai' request for model 'gpt-99' failed.
+Model 'gpt-99' was not found. Confirm the model identifier is correct and your openai account has access to it.
+```
+
+Resolution -- verify the model name and that your account has access. Use the
+`provider/model` format, e.g. `openai/gpt-5-mini`, `anthropic/claude-sonnet-4-5`.
+
+## Test Mode Errors
+
+These errors occur when running `osmosis test` or `osmosis eval`.
+
+### DatasetParseError
+
+Raised when the dataset file cannot be read or parsed.
+
+```
+DatasetParseError: Unsupported file format: .txt. Supported formats: .parquet (recommended), .jsonl, .csv
+```
+
+```
+DatasetParseError: Invalid JSON at line 5: Expecting ',' delimiter
+```
+
+```
+DatasetParseError: Parquet support requires pyarrow. Install with: pip install pyarrow
+```
+
+Resolution:
+
+- Ensure your file uses a supported format: `.parquet`, `.jsonl`, or `.csv`.
+- For JSONL files, verify each line is valid JSON.
+- For Parquet files, install `pyarrow` (included in the `server` extra).
+
+### DatasetValidationError
+
+Raised when dataset rows are missing required columns or have invalid values.
+
+```
+DatasetValidationError: Row 0: Missing required columns: ['user_prompt']
+```
+
+```
+DatasetValidationError: Row 3: 'ground_truth' cannot be null
+```
+
+```
+DatasetValidationError: Row 7: 'system_prompt' must be a string, got int
+```
+
+Every dataset row must include these columns:
+
+| Column | Type | Description |
+|--------|------|-------------|
+| `ground_truth` | `str` | Reference answer |
+| `user_prompt` | `str` | User message |
+| `system_prompt` | `str` | System prompt |
+
+All values must be non-empty strings.
+
+### ToolValidationError
+
+Raised when tool schemas returned by your agent are invalid for the provider's
+API.
+
+```
+ToolValidationError:
+```
+
+Resolution -- ensure your `get_tools()` method returns valid OpenAI-compatible
+function tool schemas. Each tool must have a `type` field set to `"function"` and
+a `function` object with at least a `name`.
+
+### Provider connection errors
+
+`SystemicProviderError` is raised when a provider error affects all rows (e.g.
+authentication failure, budget exhausted, network unreachable). The batch aborts
+early instead of retrying each row.
+
+Resolution -- fix the underlying credential or connectivity issue and re-run.
+
+## Remote Rollout Errors
+
+### AgentLoopValidationError
+
+Raised by `osmosis validate` or when starting a server with `validate=True`
+(the default). The error lists all validation failures.
+
+```
+AgentLoopValidationError: Agent loop validation failed with 2 error(s):
+ - [MISSING_NAME] name: Agent loop must have a 'name' attribute
+ - [RUN_NOT_ASYNC] run: 'run' method must be an async function (async def)
+```
+
+Common validation error codes:
+
+| Code | Meaning |
+|------|---------|
+| `MISSING_NAME` | Agent loop class has no `name` attribute. |
+| `INVALID_NAME_TYPE` | `name` is not a string. |
+| `EMPTY_NAME` | `name` is empty or whitespace. |
+| `MISSING_RUN_METHOD` | No `run` method defined. |
+| `RUN_NOT_CALLABLE` | `run` exists but is not callable. |
+| `RUN_NOT_ASYNC` | `run` is not an `async def` function. |
+| `GET_TOOLS_RETURNS_NONE` | `get_tools()` returned `None` instead of a list. |
+| `GET_TOOLS_INVALID_TYPE` | `get_tools()` returned a non-list type. |
+| `GET_TOOLS_EXCEPTION` | `get_tools()` raised an exception. |
+| `MISSING_TOOL_TYPE` | A tool dict is missing the `type` field. |
+| `MISSING_FUNCTION` | A tool dict is missing the `function` field. |
+| `MISSING_FUNCTION_NAME` | A function definition has no `name`. |
+| `INVALID_FUNCTION_NAME` | Function name is empty or not a string. |
+
+Run validation before serving to catch these early:
+
+```bash
+osmosis validate -m my_agent:agent_loop
+```
+
+### ServeError
+
+Raised when `serve_agent_loop()` or `osmosis serve` cannot start the server.
+
+**Not logged in:**
+
+```
+ServeError: Not logged in. Please run 'osmosis login' first, or use skip_register=True for local testing.
+```
+
+Resolution -- run `osmosis login`, or pass `--skip-register` / `--local` if you
+do not need platform registration.
+
+**Conflicting options:**
+
+```
+ServeError: local_debug=True disables API key authentication; do not provide api_key in local debug mode.
+```
+
+Resolution -- do not combine `--local` with `--api-key`. Local debug mode
+intentionally disables API key authentication.
+
+**Missing dependencies:**
+
+```
+ImportError: FastAPI is required for serve_agent_loop(). Install it with: pip install osmosis-ai[server]
+```
+
+```
+ImportError: uvicorn is required for serve_agent_loop(). Install it with: pip install osmosis-ai[server]
+```
+
+Resolution:
+
+```bash
+pip install osmosis-ai[server]
+```
+
+### Connection / timeout issues
+
+Rollout protocol exceptions are defined in
+`osmosis_ai.rollout.core.exceptions`:
+
+| Exception | Description | Retryable? |
+|-----------|-------------|------------|
+| `OsmosisTransportError` | Network-level failure (connection refused, DNS error). | Yes |
+| `OsmosisServerError` | Server returned HTTP 5xx. Includes `status_code` attribute. | Yes |
+| `OsmosisValidationError` | Server returned HTTP 4xx. Includes `status_code` attribute. | No |
+| `OsmosisTimeoutError` | Request exceeded configured timeout. | Yes |
+| `AgentLoopNotFoundError` | Registry lookup for agent name failed. Includes `name` and `available` attributes. | No |
+| `ToolExecutionError` | A tool call failed during execution. Includes `tool_call_id` and `tool_name`. | Depends |
+| `ToolArgumentError` | Tool arguments could not be parsed (subclass of `ToolExecutionError`). | No |
+
+For retryable errors, use exponential backoff. The SDK client settings provide
+configurable retry parameters (see [Configuration](./configuration.md)).
+
+### PublicIPDetectionError
+
+Raised when the server cannot detect its public IP address. This happens when
+binding to `0.0.0.0` and all detection methods fail.
+
+```
+PublicIPDetectionError: Failed to detect public IP address. All detection methods failed:
+ 1. Cloud metadata (AWS/GCP/Azure): unavailable or returned no public IP
+ 2. External IP services: all failed or timed out
+
+To fix this, provide an explicit IP/hostname to your application.
+```
+
+Resolution -- the SDK tries cloud metadata services (AWS, GCP, Azure) first,
+then external IP services (checkip.amazonaws.com, ipify, icanhazip, ifconfig.me)
+as a fallback. If all fail:
+
+- Check network connectivity and firewall rules.
+- Provide an explicit host via the `OSMOSIS_PUBLIC_HOST` environment variable
+ or the `--host` flag.
+- If running locally, use `--local` mode which skips IP detection for
+ platform registration.
+
+## Debug Tips
+
+### Using `--verbose` flag
+
+The `osmosis serve` command accepts `-v` / `--verbose` to increase output
+verbosity, which prints detailed validation warnings and server configuration.
+
+```bash
+osmosis serve -m my_agent:agent_loop --verbose
+```
+
+### Using `--debug` flag
+
+The `osmosis test` and `osmosis eval` commands accept `--debug` to enable debug
+logging, which prints detailed information about each step including provider
+requests and responses.
+
+```bash
+osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini --debug
+osmosis eval -m my_agent:agent_loop -d data.jsonl --eval-fn rewards:fn --model openai/gpt-5-mini --debug
+```
+
+### Interactive mode in test mode
+
+Use `--interactive` with `osmosis test` to step through each row one at a time.
+This is useful for debugging agent logic and inspecting intermediate messages:
+
+```bash
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
+
+# Start at a specific row
+osmosis test -m my_agent:agent_loop -d data.jsonl --interactive --row 5
+```
+
+### Debug directory for rollout server
+
+When serving an agent, use `--debug-dir` to write detailed execution traces
+for each rollout to JSONL files:
+
+```bash
+osmosis serve -m my_agent:agent_loop --debug-dir ./debug-logs
+```
+
+Each rollout will produce a file at `{debug_dir}/{timestamp}/{rollout_id}.jsonl`.
+
+### Checking logs
+
+The SDK uses Python's standard `logging` module. Increase the log level to see
+more detail:
+
+```python
+import logging
+logging.basicConfig(level=logging.DEBUG)
+```
+
+Or set the uvicorn log level when serving:
+
+```bash
+osmosis serve -m my_agent:agent_loop --log-level debug
+```
+
+### Checking credentials
+
+To verify your authentication state and credential file:
+
+```bash
+# Show current user and workspace
+osmosis whoami
+
+# Credentials are stored at:
+# ~/.config/osmosis/credentials.json
+```
+
+## See Also
+
+- [CLI Reference](./cli.md) -- all `osmosis` commands and options
+- [Configuration](./configuration.md) -- environment variables and settings
+- [Dataset Format](./datasets.md) -- supported formats and required columns
+- [Test Mode](./test-mode.md) -- full `osmosis test` documentation
+- [Eval Mode](./eval-mode.md) -- full `osmosis eval` documentation
+- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
diff --git a/examples/README.md b/examples/README.md
index 06fd2a98..37c25244 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -26,11 +26,13 @@ All functions decorated with `@osmosis_reward` must have this exact signature:
```python
@osmosis_reward
-def your_function(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
+def your_function(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
# Your reward logic here
return float_score
```
+> **Note:** Including `**kwargs` is required for platform compatibility. The Osmosis platform passes additional keyword arguments to reward functions.
+
### Parameters
- `solution_str: str` - The solution string to evaluate
diff --git a/osmosis_ai/rollout/core/base.py b/osmosis_ai/rollout/core/base.py
index d39043d0..03b29971 100644
--- a/osmosis_ai/rollout/core/base.py
+++ b/osmosis_ai/rollout/core/base.py
@@ -124,8 +124,13 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
# Execute tools and add results
for tool_call in result.tool_calls:
- tool_result = await self.execute_tool(tool_call)
- messages.append(tool_result)
+ args = json.loads(tool_call["function"]["arguments"])
+ tool_result = do_tool(tool_call["function"]["name"], args)
+ messages.append({
+ "role": "tool",
+ "tool_call_id": tool_call["id"],
+ "content": str(tool_result),
+ })
ctx.record_tool_call(latency_ms=...)
return ctx.complete(messages)
From a8b820d1a806ee774af048f1486ce99457164f15 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 14:59:33 -0800
Subject: [PATCH 39/60] simplify README
---
README.md | 23 +++--------------------
1 file changed, 3 insertions(+), 20 deletions(-)
diff --git a/README.md b/README.md
index ce3db82e..7b6150b4 100644
--- a/README.md
+++ b/README.md
@@ -22,27 +22,10 @@ Osmosis AI is a platform for training LLMs with reinforcement learning. You defi
## Quick Start
-The fastest way to get started is with **Local Rollout** using MCP tools. Clone the example repo and test it locally in under a minute:
+Pick a training mode and follow the example repo:
-```bash
-# 1. Install the SDK with MCP support
-pip install "osmosis-ai[mcp]"
-
-# 2. Clone the example repo
-git clone https://github.com/Osmosis-AI/osmosis-git-sync-example.git
-cd osmosis-git-sync-example
-
-# 3. Set your LLM API key
-export OPENAI_API_KEY="sk-..."
-
-# 4. Run the agent against the sample dataset
-osmosis test --mcp ./mcp -d test_data.jsonl --model openai/gpt-4o-mini
-
-# 5. Try interactive mode to step through each LLM call
-osmosis test --mcp ./mcp -d test_data.jsonl --interactive
-```
-
-For detailed setup, see the [Local Rollout docs](docs/local-rollout/overview.md). For custom agent architectures, see [Remote Rollout](docs/remote-rollout/overview.md).
+- **Local Rollout** (recommended for most users): **[osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)**
+- **Remote Rollout** (custom agent architectures): **[osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)**
## Two Training Modes
From fbf4139b333c387356d731b636ffd3dd5b9c7ec4 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 15:19:36 -0800
Subject: [PATCH 40/60] Remove deployment documentation from various sections
and files to streamline content and avoid redundancy.
---
docs/README.md | 1 -
docs/configuration.md | 1 -
docs/remote-rollout/deployment.md | 109 ------------------------------
docs/remote-rollout/examples.md | 1 -
docs/remote-rollout/overview.md | 1 -
docs/remote-rollout/testing.md | 1 -
6 files changed, 114 deletions(-)
delete mode 100644 docs/remote-rollout/deployment.md
diff --git a/docs/README.md b/docs/README.md
index 97d9f668..dfbb3acf 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -25,7 +25,6 @@ Self-hosted mode where you implement and run the agent loop as a server.
- [Agent Loop Guide](./remote-rollout/agent-loop.md) -- API reference for classes, schemas, types
- [Examples](./remote-rollout/examples.md) -- agent implementations and utilities
- [Testing](./remote-rollout/testing.md) -- unit tests and mock trainer
-- [Deployment](./remote-rollout/deployment.md) -- Docker, health checks, production config
## Shared Concepts
diff --git a/docs/configuration.md b/docs/configuration.md
index 7a35d082..7394bf1d 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -230,4 +230,3 @@ OPENAI_API_KEY=sk-...
- [CLI Reference](./cli.md) -- all `osmosis` commands and options
- [Troubleshooting](./troubleshooting.md) -- common errors and resolutions
-- [Deployment](./remote-rollout/deployment.md) -- production configuration for remote rollout
diff --git a/docs/remote-rollout/deployment.md b/docs/remote-rollout/deployment.md
deleted file mode 100644
index 8e0ab565..00000000
--- a/docs/remote-rollout/deployment.md
+++ /dev/null
@@ -1,109 +0,0 @@
-# Production Deployment
-
-This guide covers deploying your rollout server to production environments using Docker, Docker Compose, health checks, and logging best practices.
-
-## Docker Configuration
-
-Using the CLI (recommended):
-
-```dockerfile
-# Dockerfile
-
-FROM python:3.11-slim
-
-WORKDIR /app
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-COPY . .
-
-EXPOSE 9000
-
-# Using CLI (validates on startup)
-CMD ["osmosis", "serve", "-m", "main:agent_loop", "-p", "9000", "--skip-register"]
-```
-
-Or with uvicorn directly:
-
-```dockerfile
-# Dockerfile (alternative)
-
-FROM python:3.11-slim
-
-WORKDIR /app
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-COPY . .
-
-EXPOSE 9000
-
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "9000"]
-```
-
-```txt
-# requirements.txt
-osmosis-ai[server,config]>=0.1.0
-```
-
-## Docker Compose
-
-```yaml
-# docker-compose.yml
-
-version: '3.8'
-
-services:
- rollout-server:
- build: .
- ports:
- - "9000:9000"
- environment:
- - OSMOSIS_ROLLOUT_LOG_LEVEL=INFO
- healthcheck:
- test: ["CMD", "curl", "-f", "http://localhost:9000/health"]
- interval: 30s
- timeout: 10s
- retries: 3
-```
-
-## Health Check Integration
-
-```python
-# For Kubernetes or load balancer health checks
-
-@app.get("/health")
-async def health():
- return {
- "status": "healthy",
- "agent_loop": agent_loop.name,
- }
-
-@app.get("/ready")
-async def ready():
- # Add any readiness checks here
- return {"ready": True}
-```
-
-## Logging Configuration
-
-```python
-import logging
-
-logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-
-# Reduce httpx noise
-logging.getLogger("httpx").setLevel(logging.WARNING)
-```
-
-## See Also
-
-- [Configuration](../configuration.md) -- environment variables and SDK settings
-- [Examples](./examples.md) -- agent implementations
-- [Testing](./testing.md) -- unit tests and mock trainer
-- [Architecture](./architecture.md) -- protocol design
diff --git a/docs/remote-rollout/examples.md b/docs/remote-rollout/examples.md
index a25ec7a1..ca2dbca4 100644
--- a/docs/remote-rollout/examples.md
+++ b/docs/remote-rollout/examples.md
@@ -691,5 +691,4 @@ For a complete, runnable project with tools, rewards, and server setup, see: [os
## See Also
- [Testing](./testing.md) -- unit tests and mock trainer
-- [Deployment](./deployment.md) -- Docker, health checks, production config
- [Agent Loop Guide](./agent-loop.md) -- endpoints, schemas, types
diff --git a/docs/remote-rollout/overview.md b/docs/remote-rollout/overview.md
index 30d2971e..19c3d4a9 100644
--- a/docs/remote-rollout/overview.md
+++ b/docs/remote-rollout/overview.md
@@ -251,6 +251,5 @@ See the complete working example: [osmosis-remote-rollout-example](https://githu
- [Agent Loop Guide](./agent-loop.md) -- Complete API documentation
- [Examples](./examples.md) -- Full working examples
- [Testing](./testing.md) -- Unit tests and mock trainer
-- [Deployment](./deployment.md) -- Docker, health checks, production config
- [Dataset Format](../datasets.md) -- Supported formats and required columns
- [Test Mode](../test-mode.md) -- Local testing with external LLM providers
diff --git a/docs/remote-rollout/testing.md b/docs/remote-rollout/testing.md
index b076efa2..429f4574 100644
--- a/docs/remote-rollout/testing.md
+++ b/docs/remote-rollout/testing.md
@@ -201,4 +201,3 @@ Generates deterministic fake prompt token IDs for testing. The token count scale
- [Examples](./examples.md) -- agent implementations and utilities
- [Test Mode](../test-mode.md) -- testing with cloud LLMs
-- [Deployment](./deployment.md) -- production deployment
From e8001e237176b303796e62d486ae5e1b53bdda5c Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 15:37:21 -0800
Subject: [PATCH 41/60] Update documentation for agent loop implementation and
enhance CSV dataset parsing
- Clarified the required methods for implementing the agent loop in the remote rollout documentation.
- Improved error handling in the DatasetReader class to check for mismatched columns in CSV files, raising a DatasetParseError for better debugging.
---
docs/remote-rollout/overview.md | 2 +-
osmosis_ai/rollout/eval/common/dataset.py | 9 +++++++--
2 files changed, 8 insertions(+), 3 deletions(-)
diff --git a/docs/remote-rollout/overview.md b/docs/remote-rollout/overview.md
index 19c3d4a9..bc29cfa0 100644
--- a/docs/remote-rollout/overview.md
+++ b/docs/remote-rollout/overview.md
@@ -30,7 +30,7 @@ pip install osmosis-ai[full]
### Step 1: Implement Your Agent Loop
-Create a class that inherits from `RolloutAgentLoop` and implements two methods:
+Create a class that inherits from `RolloutAgentLoop` and implements the required `get_tools` and `run` methods:
```python
import json
diff --git a/osmosis_ai/rollout/eval/common/dataset.py b/osmosis_ai/rollout/eval/common/dataset.py
index 6d2bc39f..51edd002 100644
--- a/osmosis_ai/rollout/eval/common/dataset.py
+++ b/osmosis_ai/rollout/eval/common/dataset.py
@@ -168,8 +168,13 @@ def _iter_csv(self) -> Iterator[dict[str, Any]]:
try:
with open(self.file_path, encoding="utf-8", newline="") as f:
reader = csv.DictReader(f)
- for row in reader:
- yield dict(row)
+ for row_num, row in enumerate(reader, start=1):
+ row_dict = dict(row)
+ if None in row_dict:
+ raise DatasetParseError(
+ f"CSV row {row_num} has more columns than the header"
+ )
+ yield row_dict
except csv.Error as e:
raise DatasetParseError(f"Invalid CSV file: {e}") from e
except OSError as e:
From c78730ee795eecdbbaacd823b2707bd98101fbac Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 16:04:11 -0800
Subject: [PATCH 42/60] simplify docs
---
docs/cli.md | 94 +---
docs/configuration.md | 169 +-----
docs/eval-mode.md | 521 ++-----------------
docs/local-rollout/reward-functions.md | 70 +--
docs/remote-rollout/agent-loop.md | 407 ++-------------
docs/remote-rollout/examples.md | 693 ++-----------------------
docs/remote-rollout/overview.md | 166 ++----
docs/remote-rollout/testing.md | 104 +---
docs/test-mode.md | 393 ++------------
docs/troubleshooting.md | 22 +-
10 files changed, 263 insertions(+), 2376 deletions(-)
diff --git a/docs/cli.md b/docs/cli.md
index e0158402..d2380954 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -66,17 +66,9 @@ osmosis workspace switch
Test your agent locally against a dataset using external LLMs. Works with both Local Rollout (MCP tools) and Remote Rollout (RolloutAgentLoop) agents.
```bash
-# Remote Rollout: test an AgentLoop implementation
-osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini
-
-# Local Rollout: test MCP tools directly (no AgentLoop needed)
+osmosis test -m server:agent_loop -d data.jsonl --model openai/gpt-5-mini
osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-
-# Interactive step-by-step debugging
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
-
-# Interactive with MCP tools
-osmosis test --mcp ./mcp -d data.jsonl --interactive
+osmosis test -m server:agent_loop -d data.jsonl --interactive
```
See [Test Mode](./test-mode.md) for full documentation on dataset format, interactive mode, and all options.
@@ -88,20 +80,12 @@ See [Test Mode](./test-mode.md) for full documentation on dataset format, intera
Evaluate trained models with custom eval functions and pass@k metrics. Works with both Local Rollout (MCP tools) and Remote Rollout (RolloutAgentLoop) agents.
```bash
-# Remote Rollout: benchmark a trained model at a serving endpoint
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
+osmosis eval -m server:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1
-# Local Rollout: evaluate MCP tools directly
osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model openai/gpt-5-mini
-
-# pass@k with 5 runs per row
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
+ --eval-fn rewards:compute_reward --model openai/gpt-5-mini
```
See [Eval Mode](./eval-mode.md) for full documentation on eval functions, pass@k metrics, and output formats.
@@ -114,30 +98,16 @@ Start a RolloutServer for an agent loop implementation. The module path format i
```bash
# Start server with Platform registration (requires `osmosis login`)
-osmosis login
-osmosis serve -m my_agent:agent_loop
-
-# Specify host and port
-osmosis serve -m my_agent:agent_loop -H 127.0.0.1 -p 8080
-
-# Local / container mode: skip Platform registration (no login required).
-# NOTE: API key auth is still enabled by default.
-osmosis serve -m my_agent:agent_loop --skip-register
+osmosis serve -m server:agent_loop
# Local debug mode: disable API key auth AND skip Platform registration
-osmosis serve -m my_agent:agent_loop --local
-
-# Provide a stable API key (otherwise one is generated and printed on startup)
-osmosis serve -m my_agent:agent_loop --skip-register --api-key "$MY_API_KEY"
+osmosis serve -m server:agent_loop --local
# Enable auto-reload for development
-osmosis serve -m my_agent:agent_loop --local --reload
+osmosis serve -m server:agent_loop --local --reload
# Set log level and write execution traces to a directory
-osmosis serve -m my_agent:agent_loop --log-level debug --log ./traces
-
-# Skip validation (not recommended)
-osmosis serve -m my_agent:agent_loop --no-validate
+osmosis serve -m server:agent_loop --log-level debug --log ./traces
```
**Options:**
@@ -153,77 +123,55 @@ osmosis serve -m my_agent:agent_loop --no-validate
| `--no-validate` | Skip agent loop validation before starting |
| `--reload` | Enable auto-reload for development |
| `--log-level` | Uvicorn log level: debug, info, warning, error, critical (default: info) |
-| `--log DIR` | Enable logging and write per-rollout execution traces (as `{rollout_id}.jsonl`) to DIR |
+| `--log DIR` | Enable logging and write per-rollout execution traces to DIR |
-> **Note:** The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This key is **not** the same as your `osmosis login` token (which is for authenticating with the Osmosis Platform), nor is it used for callbacks *from* your server back to TrainGate.
+> **Note:** The `--api-key` option sets the API key for this RolloutServer. It is used by TrainGate to authenticate its requests *to* your server. This is **not** the same as your `osmosis login` token.
See [Remote Rollout Overview](./remote-rollout/overview.md) for architecture details and the full agent lifecycle.
### osmosis validate
-Validate an agent loop before starting the server (checks tools, async run method, etc.):
+Validate an agent loop before starting the server:
```bash
-# Basic validation
-osmosis validate -m my_agent:agent_loop
-
-# Verbose output with detailed warnings
-osmosis validate -m my_agent:agent_loop -v
+osmosis validate -m server:agent_loop
+osmosis validate -m server:agent_loop -v # Verbose with warnings
```
-**Options:**
-
| Option | Description |
|--------|-------------|
| `-m`/`--module` | Module path to the agent loop (required) |
| `-v`/`--verbose` | Show detailed validation output including warnings |
-See [Remote Rollout Overview](./remote-rollout/overview.md) for architecture details and the full agent lifecycle.
-
## Rubric Tools
### osmosis preview
-Preview a rubric file and print every configuration discovered, including nested entries:
+Preview a rubric file or dataset:
```bash
osmosis preview --path path/to/rubric.yaml
-```
-
-Preview a dataset of rubric-scored solutions stored as JSONL:
-
-```bash
osmosis preview --path path/to/data.jsonl
```
-Both formats validate the file, echo a short summary (`Loaded ...`), 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.
+Both formats validate the file, echo a short summary, and pretty-print the parsed records.
### osmosis eval-rubric
-Evaluate a dataset against a hosted rubric configuration and print the returned scores:
+Evaluate a dataset against a hosted rubric configuration:
```bash
osmosis eval-rubric --rubric support_followup --data examples/sample_data.jsonl
```
-**Command split** (development-stage breaking change):
-- `osmosis eval-rubric` evaluates JSONL conversations against hosted rubrics.
-- `osmosis eval` runs rollout eval functions against agent datasets.
-
**Options:**
-- `-d`/`--data path/to/data.jsonl` -- Supply the dataset; the path is resolved relative to the current working directory.
-- `--config path/to/rubric_configs.yaml` -- Provide rubric definitions when they are not located alongside the dataset.
-- `-n`/`--number` -- Sample the provider multiple times per record; the CLI prints every run along with aggregate statistics (average, variance, standard deviation, and min/max).
-- `--output path/to/dir` -- Create the directory (if needed) and emit `rubric_eval_result_.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.
-
-**Behavior notes:**
+- `-d`/`--data path/to/data.jsonl` -- Supply the dataset
+- `--config path/to/rubric_configs.yaml` -- Provide rubric definitions
+- `-n`/`--number` -- Sample multiple times per record with aggregate statistics
+- `--output path/to/dir` -- Write results to a file or directory
-- Skip `--output` to collect results under `~/.cache/osmosis/eval_result//rubric_eval_result_.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.
-- Each dataset record must provide a non-empty `solution_str`; optional fields such as `original_input`, `ground_truth`, and `extra_info` travel with the record and are forwarded to the evaluator when present.
-- When delegating to a custom `@osmosis_rubric` function, the CLI enriches `extra_info` with the active `provider`, `model`, `rubric`, score bounds, any configured `system_prompt`, the resolved `original_input`, and the record's metadata/extra fields so the decorator's required entries are always present.
-- Rubric configuration files intentionally reject `extra_info`; provide per-example context through the dataset instead.
+**Command split** (development-stage breaking change): `osmosis eval-rubric` evaluates JSONL conversations against hosted rubrics, while `osmosis eval` runs rollout eval functions against agent datasets.
## See Also
diff --git a/docs/configuration.md b/docs/configuration.md
index 7394bf1d..df8ba6d7 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -30,194 +30,55 @@ Server-side configuration for the RolloutServer process.
| `OSMOSIS_ROLLOUT_SERVER_RECORD_TTL_SECONDS` | `float` | `3600.0` | 60.0 -- 86400.0 | How long to keep completed rollout records (seconds). |
| `OSMOSIS_ROLLOUT_SERVER_CLEANUP_INTERVAL_SECONDS` | `float` | `60.0` | 10.0 -- 3600.0 | Interval for the cleanup task (seconds). |
| `OSMOSIS_ROLLOUT_SERVER_REQUEST_TIMEOUT_SECONDS` | `float` | `600.0` | 10.0 -- 3600.0 | Timeout for individual requests (seconds). |
-| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_READINESS_TIMEOUT_SECONDS` | `float` | `10.0` | 1.0 -- 60.0 | Maximum wait for the server to become ready before platform registration. The server polls its own health endpoint to confirm readiness. |
+| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_READINESS_TIMEOUT_SECONDS` | `float` | `10.0` | 1.0 -- 60.0 | Maximum wait for the server to become ready before platform registration. |
| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_READINESS_POLL_INTERVAL_SECONDS` | `float` | `0.2` | 0.05 -- 5.0 | Interval between health check polls during server readiness check. |
| `OSMOSIS_ROLLOUT_SERVER_REGISTRATION_SHUTDOWN_TIMEOUT_SECONDS` | `float` | `30.0` | 1.0 -- 300.0 | Timeout for waiting for platform registration to complete during shutdown. |
### Global Settings (`OSMOSIS_ROLLOUT_*`)
-Top-level settings that apply across the SDK.
-
| Variable | Type | Default | Range | Description |
|----------|------|---------|-------|-------------|
| `OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES` | `int` | `1048576` (1 MB) | 1024 -- 104857600 (100 MB) | Maximum allowed size for rollout metadata in bytes. |
-## Settings Classes
-
-Configuration is organized into three Pydantic Settings classes. When `pydantic-settings` is installed, these classes automatically read from environment variables and `.env` files. Without `pydantic-settings`, they fall back to plain `pydantic.BaseModel` with only programmatic configuration.
-
-### `RolloutClientSettings`
-
-HTTP client settings for communication with TrainGate.
-
-```python
-from osmosis_ai.rollout.config import RolloutClientSettings
-
-client = RolloutClientSettings()
-print(client.timeout_seconds) # 300.0
-print(client.max_retries) # 3
-print(client.retry_base_delay) # 1.0
-```
-
-Env prefix: `OSMOSIS_ROLLOUT_CLIENT_`
-
-### `RolloutServerSettings`
-
-Server-side settings for the RolloutServer process.
-
-```python
-from osmosis_ai.rollout.config import RolloutServerSettings
-
-server = RolloutServerSettings()
-print(server.max_concurrent_rollouts) # 100
-print(server.request_timeout_seconds) # 600.0
-```
-
-Env prefix: `OSMOSIS_ROLLOUT_SERVER_`
-
-### `RolloutSettings`
-
-Top-level settings that aggregates `RolloutClientSettings` and `RolloutServerSettings` as nested objects.
-
-```python
-from osmosis_ai.rollout.config import RolloutSettings
-
-settings = RolloutSettings()
-print(settings.client.timeout_seconds) # 300.0
-print(settings.server.max_concurrent_rollouts) # 100
-print(settings.max_metadata_size_bytes) # 1048576
-```
-
-Env prefix: `OSMOSIS_ROLLOUT_`
-
## Programmatic Configuration
-### `get_settings()`
+Configuration is organized into three Pydantic Settings classes that automatically read from environment variables and `.env` files:
-Returns the global settings singleton. On first call, settings are loaded from
-environment variables and `.env` files.
+| Class | Env Prefix | Description |
+|-------|-----------|-------------|
+| `RolloutClientSettings` | `OSMOSIS_ROLLOUT_CLIENT_` | HTTP client settings (timeout, retries, connection pool) |
+| `RolloutServerSettings` | `OSMOSIS_ROLLOUT_SERVER_` | Server settings (concurrency, TTL, timeouts) |
+| `RolloutSettings` | `OSMOSIS_ROLLOUT_` | Top-level aggregator with nested `client` and `server` objects |
```python
-from osmosis_ai.rollout.config import get_settings
+from osmosis_ai.rollout.config import get_settings, configure, RolloutSettings, RolloutClientSettings
+# Read current settings (loaded from env on first call)
settings = get_settings()
-timeout = settings.client.timeout_seconds
-```
-
-### `configure()`
-
-Replaces the global settings singleton with a custom instance. Call this early
-in your application to override environment variable defaults.
-
-```python
-from osmosis_ai.rollout.config import (
- configure,
- RolloutSettings,
- RolloutClientSettings,
- RolloutServerSettings,
-)
+print(settings.client.timeout_seconds) # 300.0
+# Override settings programmatically
configure(RolloutSettings(
- client=RolloutClientSettings(
- timeout_seconds=120.0,
- max_retries=5,
- ),
- server=RolloutServerSettings(
- max_concurrent_rollouts=200,
- ),
+ client=RolloutClientSettings(timeout_seconds=120.0, max_retries=5),
))
```
-### `reset_settings()`
-
-Resets the global settings singleton to `None`. Primarily used in tests to
-ensure a clean state.
-
-```python
-from osmosis_ai.rollout.config import reset_settings
-
-reset_settings()
-```
-
-## Metadata Size Limits
-
-Rollout requests include a `metadata` dict that is validated against a
-configurable size limit. The default limit is **1 MB**. The size is measured as
-the byte length of the JSON-serialized metadata.
-
-### `get_max_metadata_size_bytes()`
-
-Returns the current maximum metadata size limit in bytes. Thread-safe.
-
-```python
-from osmosis_ai.rollout.core.schemas import get_max_metadata_size_bytes
-
-limit = get_max_metadata_size_bytes() # 1048576 (1 MB)
-```
-
-### `set_max_metadata_size_bytes()`
-
-Sets the maximum metadata size limit in bytes. Thread-safe. The value must be
-positive; otherwise a `ValueError` is raised.
-
-```python
-from osmosis_ai.rollout.core.schemas import set_max_metadata_size_bytes
-
-# Increase to 2 MB
-set_max_metadata_size_bytes(2 * 1024 * 1024)
-```
-
-The limit is also configurable via the `max_metadata_size_bytes` field on
-`RolloutSettings` (env: `OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES`), though
-the runtime functions above provide thread-safe access used during request
-validation.
+Use `reset_settings()` in tests to clear the singleton.
## Server Constants
-The `osmosis_ai.rollout.server.serve` module defines the following defaults
-for `serve_agent_loop()`:
-
| Constant | Value | Description |
|----------|-------|-------------|
| `DEFAULT_PORT` | `9000` | Default port the RolloutServer binds to. |
| `DEFAULT_HOST` | `"0.0.0.0"` | Default host the RolloutServer binds to (all interfaces). |
-Override these when starting the server:
-
-```bash
-osmosis serve -m my_agent:agent_loop -p 8080
-```
-
-Or programmatically:
-
-```python
-from osmosis_ai.rollout.server import serve_agent_loop
-
-serve_agent_loop(my_agent, host="127.0.0.1", port=8080)
-```
-
## Credential File Location
-Authentication tokens from `osmosis login` are stored at:
-
-```
-~/.config/osmosis/credentials.json
-```
-
-The configuration directory (`~/.config/osmosis/`) is created with mode `0700`
-(owner-only access). The credentials file is written with mode `0600`
-(owner read/write only).
-
-The credentials file uses a multi-workspace format. It stores credentials for
-each workspace you have logged in to, along with the currently active workspace
-name. Use `osmosis workspace list` to see all stored workspaces and
-`osmosis workspace switch ` to change the active one.
+Authentication tokens from `osmosis login` are stored at `~/.config/osmosis/credentials.json` with owner-only permissions (`0600`). The file uses a multi-workspace format -- use `osmosis workspace list` and `osmosis workspace switch ` to manage stored workspaces.
## `.env` File Support
-Both the CLI and the settings classes load environment variables from a `.env`
-file in the current working directory via `python-dotenv`. This is useful for
-local development:
+Both the CLI and the settings classes load environment variables from a `.env` file in the current working directory via `python-dotenv`:
```bash
# .env
diff --git a/docs/eval-mode.md b/docs/eval-mode.md
index faa00e4d..f3cf335e 100644
--- a/docs/eval-mode.md
+++ b/docs/eval-mode.md
@@ -13,173 +13,37 @@ Key capabilities:
- Run multiple trials per row for pass@k analysis
- Use existing `@osmosis_reward` functions or full-context eval functions
- Get statistical summaries (mean, std, min, max) per eval function
-- Compare model quality across checkpoints or configurations
+- Compare model quality with `--baseline-model`
- **Concurrent execution** with `--batch-size` for faster benchmarks
-- Use LiteLLM providers (e.g., `openai/gpt-5-mini`) as the primary model
> **Note:** Command split (development-stage breaking change):
> `osmosis eval-rubric` is for hosted rubric evaluation, while `osmosis eval` is for agent eval mode documented here.
## Quick Start
-### Benchmarking with Remote Rollout Agent
-
-The primary use case is connecting to a model serving endpoint:
-
```bash
# Benchmark a trained model served at an endpoint
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
+osmosis eval -m server:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
- --model my-finetuned-model \
- --base-url http://localhost:8000/v1
-
-# Multiple eval functions
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:exact_match \
- --eval-fn rewards:partial_match \
- --model my-finetuned-model \
- --base-url http://localhost:8000/v1
-```
-
-### Benchmarking with Local Rollout MCP Tools
-
-If you use the **git-sync** workflow and provide MCP tools (`@mcp.tool()`) instead of a `RolloutAgentLoop`, use `--mcp` to point at your MCP directory. The SDK automatically loads all registered tools and runs a standard agent loop.
-
-```bash
-# Install MCP support
-pip install osmosis-ai[mcp]
+ --model my-finetuned-model --base-url http://localhost:8000/v1
-# Evaluate using MCP tools directory
+# Local Rollout: evaluate MCP tools directory
osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn reward_fn:compute_reward \
- --model openai/gpt-5-mini
+ --eval-fn reward_fn:compute_reward --model openai/gpt-5-mini
-# Same options work: multiple eval functions, pass@k, batch, etc.
-osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn rewards:exact_match \
- --eval-fn rewards:partial_match \
- --model my-finetuned-model --base-url http://localhost:8000/v1 \
- --n 5 --batch-size 5
-```
-
-The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) for the full folder structure and examples.
-
-> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use one or the other.
-
-### Baseline Model Comparison
-
-Compare your trained model against a baseline model. The SDK runs both models on the same dataset and reports per-model summary statistics.
-
-```bash
-# Compare trained model vs GPT-5-mini baseline
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
+# Compare trained model vs baseline
+osmosis eval -m server:agent_loop -d data.jsonl \
--eval-fn rewards:compute_reward \
--model my-finetuned-model --base-url http://localhost:8000/v1 \
--baseline-model openai/gpt-5-mini
-# Compare two serving endpoints
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-model-v2 --base-url http://localhost:8000/v1 \
- --baseline-model my-model-v1 --baseline-base-url http://localhost:8001/v1
-
-# Baseline with explicit API key
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1 \
- --baseline-model anthropic/claude-sonnet-4-5 --baseline-api-key $ANTHROPIC_API_KEY
-```
-
-When a baseline is provided, the report includes **per-model summaries** with separate mean/std/min/max for each model, so you can compare their performance side by side.
-
-You can also use LiteLLM providers (e.g., `openai/gpt-5-mini`, `anthropic/claude-sonnet-4-5`) as either the primary or baseline model. See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for supported providers.
-
-### Concurrent Execution
-
-Use `--batch-size` to run multiple requests in parallel:
-
-```bash
-# Run 5 concurrent requests
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1 \
- --batch-size 5
-
-# Combine with pass@k — 10 runs per row, 5 concurrent
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 10 --batch-size 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-```
-
-### pass@k Analysis
-
-```bash
-# pass@k with 5 runs per row
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Custom pass threshold (default is 1.0)
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 --pass-threshold 0.5 \
+# pass@k with concurrent execution
+osmosis eval -m server:agent_loop -d data.jsonl \
+ --eval-fn rewards:compute_reward --n 5 --batch-size 5 \
--model my-finetuned-model --base-url http://localhost:8000/v1
```
-### Programmatic Usage
-
-```python
-from osmosis_ai.rollout.eval.common import DatasetReader, ExternalLLMClient
-from osmosis_ai.rollout.eval.evaluation import EvalRunner, EvalFnWrapper, load_eval_fns
-
-# Load dataset
-reader = DatasetReader("./test_data.jsonl")
-rows = reader.read(limit=10)
-
-# Connect to trained model endpoint
-client = ExternalLLMClient(
- "my-finetuned-model",
- api_base="http://localhost:8000/v1",
-)
-
-# Optional: baseline model for comparison
-baseline_client = ExternalLLMClient("openai/gpt-5-mini")
-
-# Load eval functions
-eval_fns = load_eval_fns(["rewards:exact_match", "rewards:partial_match"])
-
-# Create runner (with optional baseline)
-runner = EvalRunner(
- agent_loop=MyAgentLoop(),
- llm_client=client,
- eval_fns=eval_fns,
- baseline_llm_client=baseline_client, # omit for single-model mode
-)
-
-# Run evaluation (batch_size=5 for concurrent execution)
-async with client:
- async with baseline_client:
- result = await runner.run_eval(
- rows=rows,
- n_runs=5,
- max_turns=10,
- pass_threshold=0.5,
- completion_params={"temperature": 0.7},
- batch_size=5,
- )
-
-# Print results
-for name, summary in result.eval_summaries.items():
- print(f"{name}: mean={summary.mean:.3f}, std={summary.std:.3f}")
- for k, v in summary.pass_at_k.items():
- print(f" pass@{k}: {v*100:.1f}%")
-
-# Print per-model summaries (when baseline is used)
-if result.model_summaries:
- for ms in result.model_summaries:
- print(f"\n[{ms.model_tag}] {ms.model}")
- for name, summary in ms.eval_summaries.items():
- print(f" {name}: mean={summary.mean:.3f}, std={summary.std:.3f}")
-```
+The `--mcp` directory must contain a `main.py` with a `FastMCP` instance. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md). `--mcp` and `-m` are mutually exclusive.
---
@@ -193,9 +57,8 @@ Eval mode works with any OpenAI-compatible serving endpoint via `--base-url`:
| vLLM | `http://localhost:8000/v1` |
| SGLang | `http://localhost:30000/v1` |
| Ollama | `http://localhost:11434/v1` |
-| Any OpenAI-compatible API | `http://:/v1` |
-The `--model` parameter should match the model name as registered in the serving endpoint. For example, if you deployed `my-org/my-finetuned-llama` with vLLM, use `--model my-org/my-finetuned-llama`.
+The `--model` parameter should match the model name as registered in the serving endpoint. You can also use LiteLLM providers (e.g., `openai/gpt-5-mini`, `anthropic/claude-sonnet-4-5`) as either the primary or baseline model.
---
@@ -215,12 +78,10 @@ Use when you only need the final assistant response:
```python
def exact_match(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
- """Score based on the last assistant message content."""
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
```
- First parameter must be named `solution_str`
-- `solution_str` is extracted from the last assistant message
- `extra_info` contains the full row dict (all columns)
- Compatible with functions decorated with `@osmosis_reward`
@@ -230,35 +91,19 @@ Use when you need the complete conversation history:
```python
def conversation_quality(messages: list, ground_truth: str, metadata: dict, **kwargs) -> float:
- """Score based on the full conversation."""
- # messages: full conversation history (system, user, assistant, tool messages)
- # metadata: full row dict (all columns)
assistant_messages = [m for m in messages if m["role"] == "assistant"]
return min(1.0, len(assistant_messages) / 3)
```
- First parameter must be named `messages`
-- Receives the complete message list from the agent run
- Both sync and async functions are supported
-### Async Eval Functions
-
-```python
-async def llm_judge(messages: list, ground_truth: str, metadata: dict, **kwargs) -> float:
- """Use an LLM to judge the conversation quality."""
- # async eval functions are awaited automatically
- score = await call_judge_llm(messages, ground_truth)
- return score
-```
-
---
## pass@k
When `--n` is greater than 1, eval mode runs each dataset row multiple times and computes pass@k metrics.
-**pass@k** estimates the probability that at least one of `k` randomly selected samples from `n` total runs passes (score >= threshold).
-
Formula: `pass@k = 1 - C(n-c, k) / C(n, k)`
Where:
@@ -268,12 +113,6 @@ Where:
pass@k is computed per row, then averaged across all rows.
-### Example
-
-With `--n 10 --pass-threshold 0.5`:
-- If a row has 7/10 runs passing, pass@1 = 0.7, pass@3 = 0.97
-- Final pass@k values are averages across all rows
-
---
## CLI Reference
@@ -288,37 +127,37 @@ osmosis eval [OPTIONS]
|--------|-------------|
| `-d, --dataset FILE` | Path to dataset file (.parquet recommended, .jsonl, .csv) |
| `--eval-fn MODULE:FN` | Eval function path (can be specified multiple times) |
-| `--model MODEL` | Model to benchmark (see [Model Options](#model-options)) |
+| `--model MODEL` | Model to benchmark |
### Agent Options (one required, mutually exclusive)
| Option | Description |
|--------|-------------|
-| `-m, --module MODULE` | Module path to agent loop (format: `module:attribute`). Use for Remote Rollout users who implement `RolloutAgentLoop`. |
-| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for Local Rollout (git-sync) users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
+| `-m, --module MODULE` | Module path to agent loop (format: `module:attribute`) |
+| `--mcp DIR` | Path to MCP tools directory. Requires `pip install osmosis-ai[mcp]`. |
### Model Options
| Option | Description |
|--------|-------------|
-| `--model MODEL` | Required. Model name at the serving endpoint (e.g., `my-finetuned-model`), or LiteLLM provider format for baselines (e.g., `openai/gpt-5-mini`) |
-| `--base-url URL` | Base URL for the model serving endpoint (e.g., `http://localhost:8000/v1`). Use this to connect to trained model endpoints. |
-| `--api-key KEY` | API key for the endpoint or LLM provider (or use env var) |
+| `--model MODEL` | Model name at the serving endpoint, or LiteLLM format (e.g., `openai/gpt-5-mini`) |
+| `--base-url URL` | Base URL for the model serving endpoint |
+| `--api-key KEY` | API key for the endpoint or LLM provider |
### Baseline Options (optional)
| Option | Description |
|--------|-------------|
-| `--baseline-model MODEL` | Baseline model for comparison. Runs the same evaluation with a second model and reports per-model summary statistics. |
-| `--baseline-base-url URL` | Base URL for the baseline model's API endpoint. Requires `--baseline-model`. |
-| `--baseline-api-key KEY` | API key for the baseline model provider. Requires `--baseline-model`. |
+| `--baseline-model MODEL` | Baseline model for comparison (reports per-model summary statistics) |
+| `--baseline-base-url URL` | Base URL for the baseline model's endpoint |
+| `--baseline-api-key KEY` | API key for the baseline model provider |
### Execution Options
| Option | Default | Description |
|--------|---------|-------------|
| `--n N` | `1` | Number of runs per row for pass@k |
-| `--pass-threshold FLOAT` | `1.0` | Score >= threshold counts as pass for pass@k |
+| `--pass-threshold FLOAT` | `1.0` | Score >= threshold counts as pass |
| `--max-turns N` | `10` | Maximum agent turns per run |
| `--temperature FLOAT` | - | LLM sampling temperature |
| `--max-tokens N` | - | Maximum tokens per completion |
@@ -334,47 +173,6 @@ osmosis eval [OPTIONS]
| `-q, --quiet` | Suppress progress output |
| `--debug` | Enable debug logging |
-### Examples
-
-```bash
-# Remote Rollout: benchmark trained model at an endpoint
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Local Rollout: evaluate MCP tools
-osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn reward_fn:compute_reward \
- --model openai/gpt-5-mini
-
-# Local Rollout: MCP tools with trained model endpoint
-osmosis eval --mcp ./mcp -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Multiple eval functions
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:exact_match \
- --eval-fn rewards:semantic_similarity \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# pass@5 analysis
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --n 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-
-# Baseline comparison
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward \
- --model my-finetuned-model --base-url http://localhost:8000/v1 \
- --baseline-model openai/gpt-5-mini -o results.json
-
-# Concurrent execution (5 runs at a time)
-osmosis eval -m my_agent:agent_loop -d data.jsonl \
- --eval-fn rewards:compute_reward --batch-size 5 \
- --model my-finetuned-model --base-url http://localhost:8000/v1
-```
-
---
## API Reference
@@ -387,24 +185,16 @@ Orchestrates benchmark execution with eval function scoring.
from osmosis_ai.rollout.eval.evaluation import EvalRunner
runner = EvalRunner(
- agent_loop=MyAgentLoop(),
+ agent_loop=agent,
llm_client=client,
eval_fns=eval_fns,
- debug=True,
+ baseline_llm_client=baseline_client, # optional
)
-# Single run
-result = await runner.run_single(row, row_index=0, run_index=0)
-
-# Full evaluation (batch_size for concurrent execution)
-eval_result = await runner.run_eval(
- rows=rows,
- n_runs=5,
- max_turns=10,
- pass_threshold=0.5,
+result = await runner.run_eval(
+ rows=rows, n_runs=5, max_turns=10,
+ pass_threshold=0.5, batch_size=5,
completion_params={"temperature": 0.7},
- on_progress=lambda current, total, result: print(f"[{current}/{total}]"),
- batch_size=5,
)
```
@@ -417,231 +207,45 @@ eval_result = await runner.run_eval(
| `eval_fns` | `List[EvalFnWrapper]` | required | Eval functions to apply |
| `debug` | `bool` | `False` | Enable debug logging |
| `debug_dir` | `Optional[str]` | `None` | Directory for debug logs |
-| `llm_client_factory` | `Optional[Callable]` | `None` | Factory to create LLM client instances for concurrent execution. If `None`, clones from the provided `llm_client`. |
-| `baseline_llm_client` | `Optional[ExternalLLMClient]` | `None` | Baseline LLM client for comparison mode. When provided, each row is evaluated with both the primary and baseline models. |
-| `baseline_llm_client_factory` | `Optional[Callable]` | `None` | Factory to create baseline LLM client instances for concurrent execution. If `None`, clones from the provided `baseline_llm_client`. |
+| `llm_client_factory` | `Optional[Callable]` | `None` | Factory for concurrent LLM client instances |
+| `baseline_llm_client` | `Optional[ExternalLLMClient]` | `None` | Baseline LLM client for comparison mode |
+| `baseline_llm_client_factory` | `Optional[Callable]` | `None` | Factory for concurrent baseline LLM client instances |
**Methods:**
-#### `async run_single(row, row_index, run_index, max_turns, completion_params) -> EvalRunResult`
-
-Run the agent once on a row and apply all eval functions.
-
-#### `async run_eval(rows, n_runs, max_turns, completion_params, pass_threshold, on_progress, start_index, batch_size) -> EvalResult`
+- `async run_single(row, row_index, run_index, max_turns, completion_params) -> EvalRunResult`
+- `async run_eval(rows, n_runs, max_turns, completion_params, pass_threshold, on_progress, start_index, batch_size) -> EvalResult`
-Run the full benchmark across all rows. When `batch_size > 1`, runs execute concurrently using a pool of independent LLM client instances.
-
----
-
-### EvalRunResult
-
-Result from a single agent run with eval scores.
-
-```python
-@dataclass
-class EvalRunResult:
- run_index: int # Which run (0-indexed within a row)
- success: bool # Whether agent completed successfully
- scores: Dict[str, float] # Eval function name -> score
- duration_ms: float # Execution time in milliseconds
- tokens: int # Total tokens used
- error: Optional[str] # Error message (if failed)
- model_tag: Optional[str] # "primary", "baseline", or None (single-model mode)
-```
-
----
+### Result Types
-### EvalRowResult
-
-All runs for a single dataset row.
-
-```python
-@dataclass
-class EvalRowResult:
- row_index: int # Dataset row index
- runs: List[EvalRunResult] # All runs for this row
-```
-
----
-
-### EvalResult
-
-Full benchmark result with aggregated statistics.
-
-```python
-@dataclass
-class EvalResult:
- rows: List[EvalRowResult] # Per-row results
- eval_summaries: Dict[str, EvalEvalSummary] # Per-eval statistics
- total_rows: int # Number of dataset rows
- total_runs: int # Total runs (rows * n_runs)
- total_tokens: int # Total tokens consumed
- total_duration_ms: float # Total wall time
- n_runs: int # Runs per row
- pass_threshold: float # Score threshold for pass@k
- model_summaries: Optional[List[EvalModelSummary]] # Per-model stats (comparison mode)
-```
-
----
-
-### EvalEvalSummary
-
-Summary statistics for a single eval function.
-
-```python
-@dataclass
-class EvalEvalSummary:
- mean: float # Mean score across all runs
- std: float # Standard deviation
- min: float # Minimum score
- max: float # Maximum score
- pass_at_k: Dict[int, float] # k -> pass@k value (only when n > 1)
-```
-
----
-
-### EvalModelSummary
-
-Per-model summary statistics (comparison mode only).
-
-```python
-@dataclass
-class EvalModelSummary:
- model: str # Model identifier
- model_tag: str # "primary" or "baseline"
- eval_summaries: Dict[str, EvalEvalSummary] # Per-eval statistics for this model
- total_runs: int # Number of runs for this model
- total_tokens: int # Total tokens consumed
- total_duration_ms: float # Total wall time
-```
-
----
-
-### EvalFnWrapper
-
-Normalizes eval function signatures into a unified async interface.
-
-```python
-from osmosis_ai.rollout.eval.evaluation import EvalFnWrapper
-
-wrapper = EvalFnWrapper(fn=my_eval_function, name="my_eval")
-
-# Call with normalized arguments
-score = await wrapper(messages, ground_truth, metadata)
-```
-
----
+| Type | Key Fields |
+|------|------------|
+| `EvalRunResult` | `run_index`, `success`, `scores: Dict[str, float]`, `duration_ms`, `tokens`, `error`, `model_tag` |
+| `EvalRowResult` | `row_index`, `runs: List[EvalRunResult]` |
+| `EvalResult` | `rows`, `eval_summaries`, `total_rows`, `total_runs`, `total_tokens`, `n_runs`, `pass_threshold`, `model_summaries` |
+| `EvalEvalSummary` | `mean`, `std`, `min`, `max`, `pass_at_k: Dict[int, float]` |
+| `EvalModelSummary` | `model`, `model_tag`, `eval_summaries`, `total_runs`, `total_tokens`, `total_duration_ms` |
### load_eval_fns
-Load and wrap eval functions from module paths.
-
```python
from osmosis_ai.rollout.eval.evaluation import load_eval_fns
eval_fns = load_eval_fns(["rewards:exact_match", "rewards:partial_match"])
```
-Each path must be in `module:function` format. Raises `EvalFnError` if loading fails or signatures are invalid.
+Each path must be in `module:function` format. Raises `EvalFnError` if loading fails.
---
## Output Format
-When using `--output`, results are saved as JSON:
-
-```json
-{
- "config": {
- "model": "my-finetuned-model",
- "n_runs": 5,
- "pass_threshold": 0.5,
- "eval_fns": ["rewards:exact_match", "rewards:partial_match"],
- "baseline_model": "openai/gpt-5-mini"
- },
- "summary": {
- "total_rows": 100,
- "total_runs": 500,
- "eval_fns": {
- "rewards:exact_match": {
- "mean": 0.72,
- "std": 0.45,
- "min": 0.0,
- "max": 1.0,
- "pass_at_1": 0.72,
- "pass_at_3": 0.94,
- "pass_at_5": 0.98
- },
- "rewards:partial_match": {
- "mean": 0.85,
- "std": 0.22,
- "min": 0.3,
- "max": 1.0,
- "pass_at_1": 0.85,
- "pass_at_3": 0.99,
- "pass_at_5": 1.0
- }
- },
- "total_tokens": 625000,
- "total_duration_ms": 230500
- },
- "rows": [
- {
- "row_index": 0,
- "runs": [
- {
- "run_index": 0,
- "success": true,
- "scores": {
- "rewards:exact_match": 1.0,
- "rewards:partial_match": 1.0
- },
- "duration_ms": 450,
- "tokens": 200,
- "model_tag": "primary"
- },
- {
- "run_index": 0,
- "success": true,
- "scores": {
- "rewards:exact_match": 0.0,
- "rewards:partial_match": 0.7
- },
- "duration_ms": 380,
- "tokens": 180,
- "model_tag": "baseline"
- }
- ]
- }
- ],
- "model_summaries": [
- {
- "model": "my-finetuned-model",
- "model_tag": "primary",
- "total_runs": 500,
- "total_tokens": 350000,
- "total_duration_ms": 130000,
- "eval_summaries": {
- "rewards:exact_match": { "mean": 0.82, "std": 0.38, "min": 0.0, "max": 1.0 },
- "rewards:partial_match": { "mean": 0.91, "std": 0.15, "min": 0.4, "max": 1.0 }
- }
- },
- {
- "model": "openai/gpt-5-mini",
- "model_tag": "baseline",
- "total_runs": 500,
- "total_tokens": 275000,
- "total_duration_ms": 100500,
- "eval_summaries": {
- "rewards:exact_match": { "mean": 0.62, "std": 0.49, "min": 0.0, "max": 1.0 },
- "rewards:partial_match": { "mean": 0.78, "std": 0.25, "min": 0.2, "max": 1.0 }
- }
- }
- ]
-}
-```
+When using `--output`, results are saved as JSON with this structure:
-> **Note:** The `model_tag` and `model_summaries` fields are only present when `--baseline-model` is used. In single-model mode, `model_tag` is omitted from run objects and the top-level `model_summaries` key is absent.
+- `config` -- model, n_runs, pass_threshold, eval_fns, baseline_model
+- `summary` -- total_rows, total_runs, per-eval stats (mean/std/min/max/pass_at_k), total_tokens
+- `rows[]` -- per-row results with individual run scores, duration, tokens, and `model_tag`
+- `model_summaries[]` -- per-model stats (only when `--baseline-model` is used)
---
@@ -649,19 +253,6 @@ When using `--output`, results are saved as JSON:
All exceptions are shared with [Test Mode](./test-mode.md#exceptions).
-```python
-from osmosis_ai.rollout.eval.common import (
- DatasetValidationError,
- DatasetParseError,
- ProviderError,
- ToolValidationError,
-)
-from osmosis_ai.rollout.eval.evaluation import (
- EvalFnError,
- EvalModelSummary,
-)
-```
-
| Exception | Description |
|-----------|-------------|
| `EvalFnError` | Eval function loading, signature detection, or execution error |
@@ -678,26 +269,6 @@ See [Test Mode - Environment Variables](./test-mode.md#environment-variables) fo
---
-## Best Practices
-
-1. **Start with `--n 1`**: Validate your eval functions work before running expensive pass@k analyses.
-
-2. **Use `--base-url` for trained models**: Always point to your serving endpoint rather than relying on external APIs when evaluating your own models.
-
-3. **Use multiple eval functions**: Evaluate different quality dimensions (correctness, efficiency, format compliance).
-
-4. **Choose appropriate thresholds**: `--pass-threshold 1.0` (default) requires perfect scores. Use a lower threshold like `0.5` for partial-credit eval functions.
-
-5. **Compare against baselines**: Use `--baseline-model` to run the same benchmark with a second model and get per-model summary statistics.
-
-6. **Reuse `@osmosis_reward` functions**: Existing reward functions work as eval functions in simple mode without modification.
-
-7. **Use `--batch-size` to speed up benchmarks**: Concurrent execution can significantly reduce wall time, especially with high-latency endpoints. Start with a moderate value (e.g., `5`) and increase based on endpoint capacity.
-
-8. **Save results**: Use `-o results.json` to persist detailed per-run data for further analysis.
-
----
-
## See Also
- [Test Mode](./test-mode.md) -- Test agent logic with external LLMs
diff --git a/docs/local-rollout/reward-functions.md b/docs/local-rollout/reward-functions.md
index d2c5eb63..f000d71e 100644
--- a/docs/local-rollout/reward-functions.md
+++ b/docs/local-rollout/reward-functions.md
@@ -28,9 +28,7 @@ def your_function(solution_str: str, ground_truth: str, extra_info: dict = None,
The decorator raises a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
-## Examples
-
-### Exact Match
+## Example
```python
@osmosis_reward
@@ -39,69 +37,9 @@ def exact_match(solution_str: str, ground_truth: str, extra_info: dict = None, *
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
```
-### Case-Insensitive Match with Partial Credit
-
-```python
-@osmosis_reward
-def case_insensitive_match(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
- """Case-insensitive string matching with partial credit."""
- match = solution_str.lower().strip() == ground_truth.lower().strip()
-
- if extra_info and 'partial_credit' in extra_info:
- if not match and extra_info['partial_credit']:
- len_diff = abs(len(solution_str) - len(ground_truth))
- if len_diff <= 2:
- return 0.5
-
- return 1.0 if match else 0.0
-```
-
-### Numeric Tolerance
-
-```python
-@osmosis_reward
-def numeric_tolerance(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
- """Numeric comparison with configurable tolerance."""
- try:
- solution_num = float(solution_str.strip())
- truth_num = float(ground_truth.strip())
-
- tolerance = extra_info.get('tolerance', 0.01) if extra_info else 0.01
- return 1.0 if abs(solution_num - truth_num) <= tolerance else 0.0
- except ValueError:
- return 0.0
-```
-
-### Extracting Solutions from Agent Output
-
-A common pattern for tool-use agents is extracting a final numeric answer from a markdown-formatted response:
-
-```python
-import re
-from osmosis_ai import osmosis_reward
-
-def extract_solution(solution_str):
- """Extract the first number after a #### heading."""
- solution = re.search(r'####\s*([-+]?\d*\.?\d+)', solution_str)
- if not solution:
- return None
- return solution.group(1)
-
-@osmosis_reward
-def numbers_match_reward(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
- extracted = extract_solution(solution_str)
- try:
- sol_val = float(extracted)
- except (TypeError, ValueError):
- return 0.0
-
- gt_val = float(ground_truth)
- if abs(gt_val - sol_val) < 1e-7:
- return 1.0
- return 0.0
-```
-
-This example is from the [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) repository (`reward_fn/compute_reward.py`).
+For more examples (numeric tolerance, solution extraction, partial credit), see:
+- [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) -- `rewards.py`
+- [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) -- `reward_fn/compute_reward.py`
## File Placement
diff --git a/docs/remote-rollout/agent-loop.md b/docs/remote-rollout/agent-loop.md
index f5b42742..23a7a39d 100644
--- a/docs/remote-rollout/agent-loop.md
+++ b/docs/remote-rollout/agent-loop.md
@@ -11,8 +11,8 @@ Abstract base class for agent loop implementations.
```python
from osmosis_ai.rollout import RolloutAgentLoop
-class MyAgent(RolloutAgentLoop):
- name = "my_agent" # Required class attribute
+class CalculatorAgentLoop(RolloutAgentLoop):
+ name = "calculator" # Required class attribute
def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
...
@@ -27,22 +27,7 @@ class MyAgent(RolloutAgentLoop):
|-----------|------|-------------|
| `name` | `str` | Unique identifier for this agent loop (required) |
-**Subclass Validation (`__init_subclass__`):**
-
-When you define a concrete (non-abstract) subclass of `RolloutAgentLoop`, Python automatically validates at class definition time that the `name` class attribute is defined and non-empty. If it is missing or falsy, a `TypeError` is raised immediately:
-
-```python
-# This raises TypeError at class definition time:
-class BadAgent(RolloutAgentLoop):
- # Missing 'name' attribute!
- def get_tools(self, request):
- return []
- async def run(self, ctx):
- return ctx.complete([])
-# TypeError: Agent loop class BadAgent must define a 'name' class attribute
-```
-
-Abstract subclasses (those with remaining abstract methods) are not validated, so you can create intermediate base classes without defining `name`.
+Concrete subclasses must define `name`; omitting it raises `TypeError` at class definition time. Abstract subclasses (with remaining abstract methods) are not validated, so intermediate base classes don't need `name`.
**Abstract Methods:**
@@ -68,18 +53,7 @@ Execute the agent loop.
#### `get_default_tools() -> list[OpenAIFunctionToolSchema]`
-Return the default tool list for discovery and validation purposes, without requiring a real `RolloutRequest`.
-
-Internally, this calls `get_tools()` with a synthetic request (`rollout_id="discovery"`, empty messages and params). This is used by the validation framework (`validate_agent_loop()`) and platform registration to discover the agent's tools without an active rollout.
-
-```python
-agent = MyAgent()
-tools = agent.get_default_tools()
-print(f"Agent provides {len(tools)} tools")
-```
-
-- **Returns:** list of `OpenAIFunctionToolSchema` objects
-- **Notes:** If your `get_tools()` returns different tools based on request metadata, `get_default_tools()` returns the tools for an empty/default request.
+Return the default tool list for discovery and validation purposes, without requiring a real `RolloutRequest`. Calls `get_tools()` with a synthetic request. Used by the validation framework and platform registration.
---
@@ -101,7 +75,7 @@ class RolloutContext:
|-----------|------|-------------|
| `request` | `RolloutRequest` | Original rollout request |
| `tools` | `list[OpenAIFunctionToolSchema]` | Tools returned by `get_tools()` |
-| `llm` | `LLMClientProtocol` | LLM client for chat completions. In production (served via `create_app()` / `serve_agent_loop()`), this is an `OsmosisLLMClient` that calls back to TrainGate. In test/eval mode (e.g., `osmosis test`), this may be an `ExternalLLMClient` that routes requests to an external provider (OpenAI, Anthropic, etc.) via LiteLLM. |
+| `llm` | `LLMClientProtocol` | LLM client (production: `OsmosisLLMClient` calling TrainGate; test/eval: `ExternalLLMClient` via LiteLLM) |
**Methods:**
@@ -117,13 +91,6 @@ result = await ctx.chat(messages, temperature=0.7)
Create a successful completion result.
-```python
-return ctx.complete(messages, finish_reason="stop")
-
-# With reward
-return ctx.complete(messages, finish_reason="stop", reward=1.0)
-```
-
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `final_messages` | `List[Dict[str, Any]]` | required | Final conversation messages |
@@ -134,36 +101,17 @@ return ctx.complete(messages, finish_reason="stop", reward=1.0)
Create an error result.
-```python
-return ctx.error("Tool execution failed", final_messages=messages)
-```
-
#### `record_tool_call(latency_ms=0.0) -> None`
Record a tool call for metrics.
-```python
-start = time.monotonic()
-result = execute_tool(...)
-ctx.record_tool_call(latency_ms=(time.monotonic() - start) * 1000)
-```
-
#### `log_event(event_type, **data) -> None`
Log a debug event to the rollout's JSONL file. No-op if debug logging is not enabled.
```python
-# Log before LLM call
ctx.log_event("pre_llm", turn=0, num_messages=len(messages))
-
-# Log LLM response
ctx.log_event("llm_response", turn=0, has_tool_calls=result.has_tool_calls)
-
-# Log tool results
-ctx.log_event("tool_results", turn=0, num_results=len(tool_results))
-
-# Log completion
-ctx.log_event("rollout_complete", finish_reason="stop", reward=1.0)
```
**Properties:**
@@ -203,7 +151,6 @@ from osmosis_ai.rollout import OsmosisLLMClient
async with OsmosisLLMClient(
server_url="http://trainer:8080",
rollout_id="rollout-123",
- # api_key="... optional TrainGate bearer token ...",
timeout_seconds=300.0,
max_retries=3,
) as client:
@@ -216,11 +163,11 @@ async with OsmosisLLMClient(
|-----------|------|---------|-------------|
| `server_url` | `str` | required | TrainGate base URL |
| `rollout_id` | `str` | required | Unique rollout identifier |
-| `api_key` | `Optional[str]` | `None` | Optional Bearer token attached to requests to TrainGate. In the remote-rollout protocol this value is provided by TrainGate in `RolloutRequest.api_key`, and RolloutServer forwards it on callbacks to TrainGate. |
+| `api_key` | `Optional[str]` | `None` | Bearer token for TrainGate callbacks (provided by TrainGate in `RolloutRequest.api_key`) |
| `timeout_seconds` | `float` | `300.0` | Request timeout (seconds) |
| `max_retries` | `int` | `3` | Max retries for transient errors (5xx / timeouts / transport) |
| `complete_rollout_retries` | `int` | `2` | Max retries for `/v1/rollout/completed` callback |
-| `settings` | `Optional[RolloutClientSettings]` | `None` | Advanced client settings (pool/retry delays); defaults to global settings |
+| `settings` | `Optional[RolloutClientSettings]` | `None` | Advanced client settings |
**Methods:**
@@ -234,17 +181,10 @@ Call TrainGate's `/v1/chat/completions` endpoint.
| `temperature` | `float` | `1.0` | Sampling temperature |
| `top_p` | `float` | `1.0` | Top-p sampling |
| `max_tokens` | `int` | `512` | Maximum response tokens |
-| `stop` | `Optional[Union[str, List[str]]]` | `None` | Stop sequences (string is normalized to list) |
+| `stop` | `Optional[Union[str, List[str]]]` | `None` | Stop sequences |
| `logprobs` | `bool` | `True` | Return log probabilities |
-Notes:
-- Extra/unknown keyword arguments are accepted and ignored.
-
-**Raises:**
-- `OsmosisTransportError`: Network errors
-- `OsmosisServerError`: 5xx errors (after retries exhausted)
-- `OsmosisValidationError`: 4xx errors
-- `OsmosisTimeoutError`: Request timeout
+Extra/unknown keyword arguments are accepted and ignored.
#### `async complete_rollout(status, final_messages, ...) -> None`
@@ -298,26 +238,7 @@ Start a RolloutServer with automatic validation.
```python
from osmosis_ai.rollout import serve_agent_loop
-# Start with validation (default)
serve_agent_loop(agent_loop, port=9000)
-
-# Skip validation
-serve_agent_loop(agent_loop, port=9000, validate=False)
-
-# Full options
-serve_agent_loop(
- agent_loop,
- host="0.0.0.0",
- port=9000,
- validate=True,
- log_level="info",
- reload=False,
- settings=None,
- skip_register=False,
- api_key=None,
- local_debug=False,
- debug_dir=None,
-)
```
**Parameters:**
@@ -331,14 +252,12 @@ serve_agent_loop(
| `log_level` | `str` | `"info"` | Uvicorn log level |
| `reload` | `bool` | `False` | Enable auto-reload for development |
| `settings` | `Optional[RolloutSettings]` | `None` | Configuration override |
-| `skip_register` | `bool` | `False` | Skip registering with Osmosis Platform (local testing mode) |
-| `api_key` | `Optional[str]` | `None` | RolloutServer API key used by TrainGate to call this server via `Authorization: Bearer ` (generated if not provided). NOT related to `osmosis login`. |
-| `local_debug` | `bool` | `False` | Local debug mode: disable API key auth and force `skip_register=True` (NOT for production) |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging. If set, creates a timestamped subdirectory `{debug_dir}/{timestamp}/` and writes execution traces to `{rollout_id}.jsonl` files |
+| `skip_register` | `bool` | `False` | Skip registering with Osmosis Platform |
+| `api_key` | `Optional[str]` | `None` | RolloutServer API key used by TrainGate (generated if not provided) |
+| `local_debug` | `bool` | `False` | Disable API key auth and force `skip_register=True` |
+| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging |
-**Raises:**
-- `ImportError`: If FastAPI or uvicorn not installed
-- `AgentLoopValidationError`: If validation fails
+**Raises:** `ImportError` (missing FastAPI/uvicorn), `AgentLoopValidationError` (validation fails)
---
@@ -349,27 +268,7 @@ Factory function to create a FastAPI application.
```python
from osmosis_ai.rollout import create_app
-app = create_app(
- agent_loop,
- max_concurrent=100,
- record_ttl_seconds=3600,
- debug_dir="./rollout_logs", # Optional: enable debug logging
-)
-
-# Full options
-app = create_app(
- agent_loop,
- max_concurrent=100,
- record_ttl_seconds=3600,
- settings=None,
- credentials=my_credentials, # For platform registration
- server_host="0.0.0.0",
- server_port=9000,
- api_key="my-secret-key",
- debug_dir="./rollout_logs",
- on_startup=my_startup_callback,
- on_shutdown=my_shutdown_callback,
-)
+app = create_app(agent_loop, max_concurrent=100, record_ttl_seconds=3600)
```
**Parameters:**
@@ -377,18 +276,16 @@ app = create_app(
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `agent_loop` | `RolloutAgentLoop` | required | Your agent implementation |
-| `max_concurrent` | `int \| None` | `None` | Max concurrent rollouts (defaults to settings) |
-| `record_ttl_seconds` | `float \| None` | `None` | TTL for completed records (defaults to settings) |
-| `settings` | `Optional[RolloutSettings]` | `None` | Global settings override (server/logging/tracing) |
-| `credentials` | `Optional[WorkspaceCredentials]` | `None` | Workspace credentials for platform registration. If `None`, registration is skipped |
-| `server_host` | `Optional[str]` | `None` | Host the server is bound to (used for platform registration) |
-| `server_port` | `Optional[int]` | `None` | Port the server is listening on (used for platform registration) |
-| `api_key` | `Optional[str]` | `None` | API key for authenticating incoming requests. If provided, requests must include `Authorization: Bearer ` |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging. If set, each rollout writes execution traces to `{debug_dir}/{rollout_id}.jsonl`. Note: when using `serve_agent_loop()` or CLI, a timestamped subdirectory is created automatically |
-| `on_startup` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async callback to run during application startup (e.g., warming caches, starting background services) |
-| `on_shutdown` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async callback to run during application shutdown (e.g., stopping services, releasing resources) |
-
-**Returns:** `FastAPI` application
+| `max_concurrent` | `int \| None` | `None` | Max concurrent rollouts |
+| `record_ttl_seconds` | `float \| None` | `None` | TTL for completed records |
+| `settings` | `Optional[RolloutSettings]` | `None` | Global settings override |
+| `credentials` | `Optional[WorkspaceCredentials]` | `None` | Workspace credentials for platform registration |
+| `server_host` | `Optional[str]` | `None` | Host (for platform registration) |
+| `server_port` | `Optional[int]` | `None` | Port (for platform registration) |
+| `api_key` | `Optional[str]` | `None` | API key for authenticating incoming requests |
+| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging |
+| `on_startup` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async startup callback |
+| `on_shutdown` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async shutdown callback |
**Generated Endpoints:**
@@ -396,7 +293,7 @@ app = create_app(
|----------|--------|--------|-------------|
| `/v1/rollout/init` | POST | 202 | Accept rollout request |
| `/health` | GET | 200 | Health check (unauthenticated) |
-| `/platform/health` | GET | 200 | Authenticated health check for Osmosis Platform. Requires `Authorization: Bearer `. Returns 404 if `api_key` is not configured (e.g., `local_debug` mode). Used by the platform to verify reachability and API key correctness. |
+| `/platform/health` | GET | 200 | Authenticated health check for Osmosis Platform |
---
@@ -410,15 +307,8 @@ Validate a RolloutAgentLoop implementation.
from osmosis_ai.rollout import validate_agent_loop
result = validate_agent_loop(agent_loop)
-
-if result.valid:
- print(f"Agent '{result.agent_name}' is valid with {result.tool_count} tools")
-else:
- for error in result.errors:
- print(f"Error: {error}")
-
-# Or raise exception if invalid
-result.raise_if_invalid()
+if not result.valid:
+ result.raise_if_invalid()
```
**Parameters:**
@@ -428,7 +318,7 @@ result.raise_if_invalid()
| `agent_loop` | `RolloutAgentLoop` | required | Agent loop to validate |
| `request` | `Optional[RolloutRequest]` | `None` | Custom request for testing get_tools() |
-**Returns:** `ValidationResult`
+**Returns:** `ValidationResult` with `valid`, `errors`, `warnings`, `agent_name`, `tool_count`.
**Validation Checks:**
- `name` attribute is defined and non-empty
@@ -436,130 +326,20 @@ result.raise_if_invalid()
- Each tool conforms to OpenAI function schema
- `run()` method is async
----
-
-### ValidationResult
-
-Result of agent loop validation.
-
-```python
-@dataclass
-class ValidationResult:
- valid: bool
- errors: List[ValidationError]
- warnings: List[ValidationError]
- agent_name: Optional[str]
- tool_count: int
-```
-
-**Methods:**
-
-#### `raise_if_invalid() -> None`
-
-Raise `AgentLoopValidationError` if validation failed.
-
-#### `__bool__() -> bool`
-
-Returns `True` if validation passed.
-
-```python
-result = validate_agent_loop(agent_loop)
-if result: # Same as result.valid
- print("Valid!")
-```
-
----
-
-### ValidationError
-
-A single validation error or warning.
-
-```python
-@dataclass
-class ValidationError:
- code: str # e.g., "MISSING_NAME", "INVALID_TOOL_TYPE"
- message: str # Human-readable message
- field: Optional[str] # Field that caused error
- details: Optional[Dict[str, Any]] # Additional context
-```
-
-**Common Error Codes:**
-
-| Code | Description |
-|------|-------------|
-| `MISSING_NAME` | Agent loop has no `name` attribute |
-| `EMPTY_NAME` | `name` is empty or whitespace |
-| `GET_TOOLS_EXCEPTION` | `get_tools()` raised an exception |
-| `GET_TOOLS_RETURNS_NONE` | `get_tools()` returned `None` |
-| `MISSING_TOOL_TYPE` | Tool missing `type` field |
-| `MISSING_FUNCTION` | Tool missing `function` field |
-| `MISSING_FUNCTION_NAME` | Function missing `name` field |
-| `RUN_NOT_ASYNC` | `run()` method is not async |
-
----
-
-### AgentLoopValidationError
-
-Exception raised when validation fails.
-
-```python
-from osmosis_ai.rollout import AgentLoopValidationError
-
-try:
- result.raise_if_invalid()
-except AgentLoopValidationError as e:
- print(f"Validation failed: {e}")
- for error in e.errors:
- print(f" - {error}")
-```
+**Common Error Codes:** `MISSING_NAME`, `EMPTY_NAME`, `GET_TOOLS_EXCEPTION`, `GET_TOOLS_RETURNS_NONE`, `MISSING_TOOL_TYPE`, `MISSING_FUNCTION`, `MISSING_FUNCTION_NAME`, `RUN_NOT_ASYNC`
---
## Registry
-### register_agent_loop()
-
-Register an agent loop with the global registry.
-
-```python
-from osmosis_ai.rollout import register_agent_loop
-
-register_agent_loop(MyAgentLoop())
-```
-
-**Raises:** `ValueError` if name already registered
-
-### get_agent_loop()
+Functions for managing multiple agent loop instances:
-Get an agent loop from the global registry.
-
-```python
-from osmosis_ai.rollout import get_agent_loop
-
-loop = get_agent_loop("my_agent")
-```
-
-**Raises:** `AgentLoopNotFoundError` if not found
-
-### list_agent_loops()
-
-List all registered agent loop names.
-
-```python
-from osmosis_ai.rollout import list_agent_loops
-
-names = list_agent_loops() # ["agent1", "agent2"]
-```
-
-### unregister_agent_loop()
-
-Remove an agent loop from the registry.
-
-```python
-from osmosis_ai.rollout import unregister_agent_loop
-
-success = unregister_agent_loop("my_agent") # True or False
-```
+| Function | Description |
+|----------|-------------|
+| `register_agent_loop(instance)` | Register an agent loop. Raises `ValueError` if name is already registered. |
+| `get_agent_loop(name)` | Get by name. Raises `AgentLoopNotFoundError` if not found. |
+| `list_agent_loops()` | List all registered agent loop names. |
+| `unregister_agent_loop(name)` | Remove from registry. Returns `True` if found. |
---
@@ -579,7 +359,7 @@ class RolloutRequest(BaseModel):
max_turns: int = 10
max_tokens_total: int = 8192
metadata: Dict[str, Any] = {} # JSON-serializable (size-limited; default 1MB)
- api_key: Optional[str] = None # Optional Bearer token RolloutServer uses for callbacks to TrainGate
+ api_key: Optional[str] = None # Optional Bearer token for callbacks to TrainGate
idempotency_key: Optional[str] = None
```
@@ -617,40 +397,18 @@ OpenAI-compatible tool definition.
class OpenAIFunctionToolSchema(BaseModel):
type: str # "function"
function: OpenAIFunctionSchema
-```
-### OpenAIFunctionSchema
-
-Function definition within a tool.
-
-```python
class OpenAIFunctionSchema(BaseModel):
name: str
description: str
- parameters: OpenAIFunctionParametersSchema = OpenAIFunctionParametersSchema(
- type="object",
- properties={},
- required=[],
- )
+ parameters: OpenAIFunctionParametersSchema # defaults to empty object
strict: bool = False
-```
-
-### OpenAIFunctionParametersSchema
-JSON Schema for function parameters.
-
-```python
class OpenAIFunctionParametersSchema(BaseModel):
type: str # usually "object"
properties: Dict[str, OpenAIFunctionPropertySchema]
required: List[str]
-```
-
-### OpenAIFunctionPropertySchema
-Single parameter definition.
-
-```python
class OpenAIFunctionPropertySchema(BaseModel):
type: str
description: Optional[str] = None
@@ -675,8 +433,6 @@ class RolloutMetrics(BaseModel):
### RolloutStatus
-Enum for rollout status.
-
```python
class RolloutStatus(str, Enum):
COMPLETED = "COMPLETED"
@@ -689,80 +445,17 @@ class RolloutStatus(str, Enum):
All exceptions inherit from `OsmosisRolloutError`.
-### OsmosisRolloutError
-
-Base exception for all rollout errors.
-
-```python
-try:
- await client.chat_completions(messages)
-except OsmosisRolloutError as e:
- print(f"Rollout error: {e}")
-```
-
-### OsmosisTransportError
-
-Network/transport level errors.
-
-```python
-# Connection failed, DNS resolution failed, etc.
-```
-
-### OsmosisServerError
-
-Server returned 5xx error (retryable).
-
-```python
-except OsmosisServerError as e:
- print(f"Server error {e.status_code}: {e}")
-```
-
-### OsmosisValidationError
-
-Server returned 4xx error (not retryable).
-
-```python
-except OsmosisValidationError as e:
- print(f"Validation error {e.status_code}: {e}")
-```
-
-### OsmosisTimeoutError
-
-Request timed out.
-
-```python
-except OsmosisTimeoutError as e:
- print(f"Timeout: {e}")
-```
-
-### AgentLoopNotFoundError
-
-Agent loop not found in registry.
-
-```python
-except AgentLoopNotFoundError as e:
- print(f"Agent '{e.name}' not found. Available: {e.available}")
-```
-
-### ToolExecutionError
-
-Tool execution failed (optionally includes `tool_call_id` and `tool_name`).
-
-```python
-from osmosis_ai.rollout import ToolExecutionError
-
-raise ToolExecutionError("Tool failed", tool_call_id="call_123", tool_name="add")
-```
-
-### ToolArgumentError
-
-Tool argument parsing failed (inherits from `ToolExecutionError`).
-
-```python
-from osmosis_ai.rollout import ToolArgumentError
-
-raise ToolArgumentError("Invalid JSON", tool_call_id="call_123", tool_name="add")
-```
+| Exception | Description |
+|-----------|-------------|
+| `OsmosisRolloutError` | Base exception for all rollout errors |
+| `OsmosisTransportError` | Network/transport level errors |
+| `OsmosisServerError` | Server returned 5xx (retryable). Has `status_code` attribute. |
+| `OsmosisValidationError` | Server returned 4xx (not retryable). Has `status_code` attribute. |
+| `OsmosisTimeoutError` | Request timed out |
+| `AgentLoopNotFoundError` | Agent loop not found in registry. Has `name` and `available` attributes. |
+| `AgentLoopValidationError` | Validation failed. Has `errors` attribute. |
+| `ToolExecutionError` | Tool execution failed. Has optional `tool_call_id` and `tool_name`. |
+| `ToolArgumentError` | Tool argument parsing failed (inherits from `ToolExecutionError`). |
---
diff --git a/docs/remote-rollout/examples.md b/docs/remote-rollout/examples.md
index ca2dbca4..6c69f4e7 100644
--- a/docs/remote-rollout/examples.md
+++ b/docs/remote-rollout/examples.md
@@ -2,692 +2,95 @@
Complete working examples for the Osmosis Remote Rollout SDK.
-## CLI Quick Start
+
+Full source code is available in the [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) repository. The snippets below are abbreviated for clarity.
+
-The fastest way to run a RolloutServer is using the CLI:
+## CLI Quick Start
```bash
# Validate agent loop (checks tools, async run, etc.)
-osmosis validate -m my_agent:agent_loop
-
-# Start server with validation (default port 9000)
-osmosis serve -m my_agent:agent_loop
+osmosis validate -m server:agent_loop
-# Specify port
-osmosis serve -m my_agent:agent_loop -p 8080
+# Start server (default port 9000)
+osmosis serve -m server:agent_loop
-# Skip validation (not recommended)
-osmosis serve -m my_agent:agent_loop --no-validate
+# Start with auto-reload and debug logging
+osmosis serve -m server:agent_loop --reload --log ./logs
-# Enable auto-reload for development
-osmosis serve -m my_agent:agent_loop --reload
+# Test locally with a cloud LLM
+osmosis test -m server:agent_loop -d test_data.jsonl --model gpt-5-mini
-# Enable debug logging (writes {rollout_id}.jsonl files to ./logs/)
-osmosis serve -m my_agent:agent_loop --log ./logs
-
-# Verbose validation output
-osmosis validate -m my_agent:agent_loop -v
+# Evaluate with reward function
+osmosis eval -m server:agent_loop -d test_data.jsonl \
+ --eval-fn rewards:compute_reward --model gpt-5-mini
```
The module path format is `module:attribute`. The CLI automatically adds the current directory to Python path.
---
-## Basic Calculator Agent
-
-A simple agent that can perform arithmetic operations.
-
-```python
-# calculator_agent.py
-
-import json
-from typing import List
-
-from osmosis_ai.rollout import (
- RolloutAgentLoop,
- RolloutContext,
- RolloutRequest,
- RolloutResult,
- OpenAIFunctionToolSchema,
- OpenAIFunctionSchema,
- OpenAIFunctionParametersSchema,
- OpenAIFunctionPropertySchema,
- create_app,
-)
-
-
-class CalculatorAgent(RolloutAgentLoop):
- """Agent that can perform basic arithmetic."""
-
- name = "calculator"
-
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
- """Define the calculator tool."""
- return [
- OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="calculate",
- description="Perform arithmetic calculation",
- parameters=OpenAIFunctionParametersSchema(
- properties={
- "operation": OpenAIFunctionPropertySchema(
- type="string",
- description="The operation to perform",
- enum=["add", "subtract", "multiply", "divide"],
- ),
- "a": OpenAIFunctionPropertySchema(
- type="number",
- description="First operand",
- ),
- "b": OpenAIFunctionPropertySchema(
- type="number",
- description="Second operand",
- ),
- },
- required=["operation", "a", "b"],
- ),
- ),
- ),
- ]
-
- def execute_tool(self, name: str, args: dict) -> str:
- """Execute a tool and return the result."""
- if name == "calculate":
- op = args["operation"]
- a, b = args["a"], args["b"]
-
- if op == "add":
- result = a + b
- elif op == "subtract":
- result = a - b
- elif op == "multiply":
- result = a * b
- elif op == "divide":
- if b == 0:
- return "Error: Division by zero"
- result = a / b
- else:
- return f"Error: Unknown operation {op}"
-
- return str(result)
-
- return f"Error: Unknown tool {name}"
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- """Execute the agent loop."""
- messages = list(ctx.request.messages)
-
- for turn in range(ctx.request.max_turns):
- # Call LLM
- result = await ctx.chat(
- messages,
- **ctx.request.completion_params,
- )
- messages.append(result.message)
+## Project Structure
- # Check if done
- if not result.has_tool_calls:
- return ctx.complete(messages)
-
- # Process tool calls
- for tool_call in result.tool_calls:
- tool_name = tool_call["function"]["name"]
- tool_args = json.loads(tool_call["function"]["arguments"])
-
- # Execute tool
- tool_result = self.execute_tool(tool_name, tool_args)
- ctx.record_tool_call()
-
- # Add tool result to messages
- messages.append({
- "role": "tool",
- "content": tool_result,
- "tool_call_id": tool_call["id"],
- })
-
- return ctx.complete(messages, finish_reason="max_turns")
-
-
-# Export instance for CLI usage
-agent_loop = CalculatorAgent()
-
-# Create FastAPI app
-app = create_app(agent_loop)
-
-
-if __name__ == "__main__":
- import uvicorn
- uvicorn.run(app, host="0.0.0.0", port=9000)
```
-
-Run with CLI (recommended):
-```bash
-# Validate first
-osmosis validate -m calculator_agent:agent_loop
-
-# Start server
-osmosis serve -m calculator_agent:agent_loop -p 9000
-```
-
-Or with uvicorn directly:
-```bash
-uvicorn calculator_agent:app --port 9000
-```
-
----
-
-## Multi-Tool Agent
-
-An agent with multiple tools and dynamic tool selection.
-
-```python
-# multi_tool_agent.py
-
-import json
-from typing import List
-
-from osmosis_ai.rollout import (
- RolloutAgentLoop,
- RolloutContext,
- RolloutRequest,
- RolloutResult,
- OpenAIFunctionToolSchema,
- OpenAIFunctionSchema,
- OpenAIFunctionParametersSchema,
- OpenAIFunctionPropertySchema,
- create_app,
-)
-
-
-# Tool definitions
-SEARCH_TOOL = OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="search",
- description="Search for information on a topic",
- parameters=OpenAIFunctionParametersSchema(
- properties={
- "query": OpenAIFunctionPropertySchema(
- type="string",
- description="Search query",
- ),
- },
- required=["query"],
- ),
- ),
-)
-
-WEATHER_TOOL = OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="get_weather",
- description="Get current weather for a location",
- parameters=OpenAIFunctionParametersSchema(
- properties={
- "location": OpenAIFunctionPropertySchema(
- type="string",
- description="City name",
- ),
- },
- required=["location"],
- ),
- ),
-)
-
-CALCULATOR_TOOL = OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="calculate",
- description="Perform arithmetic",
- parameters=OpenAIFunctionParametersSchema(
- properties={
- "expression": OpenAIFunctionPropertySchema(
- type="string",
- description="Math expression to evaluate",
- ),
- },
- required=["expression"],
- ),
- ),
-)
-
-
-class MultiToolAgent(RolloutAgentLoop):
- """Agent with multiple tools."""
-
- name = "multi_tool_agent"
-
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
- """Return tools based on metadata configuration."""
- # Check if specific tools are requested in metadata
- requested_tools = request.metadata.get("enabled_tools")
-
- if requested_tools:
- tools = []
- tool_map = {
- "search": SEARCH_TOOL,
- "weather": WEATHER_TOOL,
- "calculator": CALCULATOR_TOOL,
- }
- for name in requested_tools:
- if name in tool_map:
- tools.append(tool_map[name])
- return tools
-
- # Default: return all tools
- return [SEARCH_TOOL, WEATHER_TOOL, CALCULATOR_TOOL]
-
- async def execute_tool(self, name: str, args: dict) -> str:
- """Execute a tool (mock implementations)."""
- if name == "search":
- return f"Search results for '{args['query']}': [Mock results]"
-
- if name == "get_weather":
- return f"Weather in {args['location']}: 72°F, Sunny"
-
- if name == "calculate":
- try:
- # WARNING: eval is unsafe in production!
- # Use a proper math parser instead
- result = eval(args["expression"])
- return str(result)
- except Exception as e:
- return f"Error: {e}"
-
- return f"Unknown tool: {name}"
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- """Execute the agent loop."""
- messages = list(ctx.request.messages)
-
- for _ in range(ctx.request.max_turns):
- result = await ctx.chat(messages, **ctx.request.completion_params)
- messages.append(result.message)
-
- if not result.has_tool_calls:
- return ctx.complete(messages)
-
- for tool_call in result.tool_calls:
- name = tool_call["function"]["name"]
- args = json.loads(tool_call["function"]["arguments"])
-
- tool_result = await self.execute_tool(name, args)
- ctx.record_tool_call()
-
- messages.append({
- "role": "tool",
- "content": tool_result,
- "tool_call_id": tool_call["id"],
- })
-
- return ctx.complete(messages, finish_reason="max_turns")
-
-
-app = create_app(MultiToolAgent())
-```
-
----
-
-## Agent with Error Handling
-
-Robust error handling in the agent loop.
-
-```python
-# robust_agent.py
-
-import json
-import logging
-from typing import List
-
-from osmosis_ai.rollout import (
- RolloutAgentLoop,
- RolloutContext,
- RolloutRequest,
- RolloutResult,
- OpenAIFunctionToolSchema,
- OsmosisRolloutError,
- create_app,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class RobustAgent(RolloutAgentLoop):
- """Agent with comprehensive error handling."""
-
- name = "robust_agent"
-
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
- return []
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- messages = list(ctx.request.messages)
-
- for turn in range(ctx.request.max_turns):
- try:
- result = await ctx.chat(
- messages,
- **ctx.request.completion_params,
- )
- messages.append(result.message)
-
- if not result.has_tool_calls:
- return ctx.complete(messages)
-
- # Process tools with individual error handling
- for tool_call in result.tool_calls:
- try:
- tool_result = await self.execute_tool_safe(tool_call)
- messages.append({
- "role": "tool",
- "content": tool_result,
- "tool_call_id": tool_call["id"],
- })
- ctx.record_tool_call()
- except Exception as e:
- logger.error(f"Tool execution failed: {e}")
- # Return error message to LLM so it can recover
- messages.append({
- "role": "tool",
- "content": f"Error: {str(e)}",
- "tool_call_id": tool_call["id"],
- })
-
- except OsmosisRolloutError as e:
- # Network/server errors - return error result
- logger.error(f"LLM call failed: {e}")
- return ctx.error(
- f"LLM call failed: {e}",
- final_messages=messages,
- )
-
- return ctx.complete(messages, finish_reason="max_turns")
-
- async def execute_tool_safe(self, tool_call: dict) -> str:
- """Execute tool with timeout and error handling."""
- import asyncio
-
- async def execute():
- name = tool_call["function"]["name"]
- args = json.loads(tool_call["function"]["arguments"])
- # Your tool execution logic here
- return f"Result for {name}"
-
- try:
- return await asyncio.wait_for(execute(), timeout=30.0)
- except asyncio.TimeoutError:
- raise RuntimeError("Tool execution timed out")
-
-
-app = create_app(RobustAgent())
+rollout-server/
+├── server.py # Agent loop + FastAPI app
+├── tools.py # Tool definitions and execution
+├── rewards.py # Reward computation
+├── test_data.jsonl # Test dataset
+└── pyproject.toml # Dependencies: osmosis-ai[server]>=0.2.14
```
---
-## Agent with Debug Logging
-
-Using `ctx.log_event()` to track execution for debugging and analysis.
-
-```python
-# logging_agent.py
-
-import json
-from typing import List
-
-from osmosis_ai.rollout import (
- RolloutAgentLoop,
- RolloutContext,
- RolloutRequest,
- RolloutResult,
- OpenAIFunctionToolSchema,
- create_app,
-)
-
-
-class LoggingAgent(RolloutAgentLoop):
- """Agent that logs execution details for debugging."""
-
- name = "logging_agent"
-
- def get_tools(self, request: RolloutRequest) -> List[OpenAIFunctionToolSchema]:
- return [] # Add your tools here
-
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- messages = list(ctx.request.messages)
-
- for turn in range(ctx.request.max_turns):
- # Log state before LLM call
- ctx.log_event(
- "pre_llm",
- turn=turn,
- num_messages=len(messages),
- messages_summary=[
- {"role": m["role"], "content_preview": str(m.get("content", ""))[:50]}
- for m in messages
- ],
- )
-
- # Call LLM
- result = await ctx.chat(messages, **ctx.request.completion_params)
- messages.append(result.message)
-
- # Log LLM response
- ctx.log_event(
- "llm_response",
- turn=turn,
- has_tool_calls=result.has_tool_calls,
- finish_reason=result.finish_reason,
- content_preview=str(result.content or "")[:100],
- )
-
- if not result.has_tool_calls:
- break
-
- # Process tool calls
- tool_results = []
- for tool_call in result.tool_calls:
- name = tool_call["function"]["name"]
- args = json.loads(tool_call["function"]["arguments"])
-
- tool_result = await self.execute_tool(name, args)
- tool_results.append({
- "role": "tool",
- "content": tool_result,
- "tool_call_id": tool_call["id"],
- })
- ctx.record_tool_call()
-
- messages.extend(tool_results)
+## Debug Logging
- # Log tool execution results
- ctx.log_event(
- "tool_results",
- turn=turn,
- num_results=len(tool_results),
- results_summary=[
- {"tool_call_id": r["tool_call_id"], "content_preview": r["content"][:50]}
- for r in tool_results
- ],
- )
+Enable debug logging with the `--log` CLI flag or `ROLLOUT_DEBUG_DIR` environment variable. Use `ctx.log_event()` in your agent's `run()` method to record events -- these are no-ops unless debug logging is enabled.
- # Log completion with reward (if computed)
- reward = self.compute_reward(messages, ctx.request.metadata)
- ctx.log_event(
- "rollout_complete",
- finish_reason="stop" if turn < ctx.request.max_turns - 1 else "max_turns",
- total_turns=turn + 1,
- final_message_count=len(messages),
- reward=reward,
- )
-
- return ctx.complete(messages, reward=reward)
-
- async def execute_tool(self, name: str, args: dict) -> str:
- # Your tool execution logic
- return f"Result for {name}"
-
- def compute_reward(self, messages: list, metadata: dict) -> float:
- # Your reward computation logic
- return 1.0
-
-
-# Run with: osmosis serve -m logging_agent:agent_loop --log ./rollout_logs
-agent_loop = LoggingAgent()
-app = create_app(agent_loop)
-```
-
-Each server session creates a timestamped subdirectory, and each rollout creates a JSONL file:
+```bash
+# Via CLI flag
+osmosis serve -m server:agent_loop --log ./rollout_logs
-```
-rollout_logs/
-├── 1703270400/ # Unix timestamp when server started
-│ ├── rollout-abc123.jsonl
-│ └── rollout-def456.jsonl
-└── 1703274000/ # Another server session
- └── rollout-xyz789.jsonl
+# Via environment variable
+ROLLOUT_DEBUG_DIR=./rollout_logs osmosis serve -m server:agent_loop
```
-Example log file contents:
-
-```jsonl
-{"event": "pre_llm", "rollout_id": "rollout-abc123", "turn": 0, "num_messages": 1, ...}
-{"event": "llm_response", "rollout_id": "rollout-abc123", "turn": 0, "has_tool_calls": true, ...}
-{"event": "tool_results", "rollout_id": "rollout-abc123", "turn": 0, "num_results": 1, ...}
-{"event": "pre_llm", "rollout_id": "rollout-abc123", "turn": 1, "num_messages": 4, ...}
-{"event": "llm_response", "rollout_id": "rollout-abc123", "turn": 1, "has_tool_calls": false, ...}
-{"event": "rollout_complete", "rollout_id": "rollout-abc123", "total_turns": 2, "reward": 1.0, ...}
-```
+Each server session creates a timestamped subdirectory, and each rollout creates a JSONL file with events like `pre_llm`, `llm_response`, `tool_results`, and `rollout_complete`.
---
## Tool Utilities
-The SDK provides utilities for creating and executing tool calls.
-
-```python
-# Using tool utilities
-
-from osmosis_ai.rollout import (
- create_tool_result,
- create_tool_error_result,
- serialize_tool_result,
- parse_tool_arguments,
- get_tool_call_info,
- execute_tool_calls,
-)
-
-# Create a standardized tool result message
-result_msg = create_tool_result("call_123", "42")
-# {"role": "tool", "content": "42", "tool_call_id": "call_123"}
-
-# Serialize various types to string
-serialize_tool_result(42) # "42"
-serialize_tool_result(3.14) # "3.14"
-serialize_tool_result({"a": 1}) # '{"a": 1}'
-
-# Parse arguments (handles both str and dict formats)
-args = parse_tool_arguments('{"a": 5, "b": 3}')
-args = parse_tool_arguments({"a": 5, "b": 3})
-
-# Extract tool call info
-tool_call = {
- "id": "call_123",
- "type": "function",
- "function": {"name": "add", "arguments": '{"a": 5, "b": 3}'},
-}
-call_id, name, args = get_tool_call_info(tool_call)
-# ("call_123", "add", {"a": 5, "b": 3})
+The `osmosis_ai.rollout.tools` module provides utilities for tool call execution:
-# Execute multiple tool calls concurrently
-async def my_executor(tool_call):
- call_id, name, args = get_tool_call_info(tool_call)
- result = await my_tools[name](**args)
- return create_tool_result(call_id, serialize_tool_result(result))
+| Function | Description |
+|----------|-------------|
+| `get_tool_call_info(tool_call)` | Extract `(call_id, name, args)` from a tool call dict, parsing JSON arguments |
+| `serialize_tool_result(value)` | Convert any value to a string suitable for tool results |
+| `create_tool_result(call_id, content)` | Create a `{"role": "tool", ...}` message dict |
+| `create_tool_error_result(call_id, error)` | Create an error tool result message |
+| `execute_tool_calls(calls, executor)` | Execute multiple tool calls concurrently via an async executor function |
-results = await execute_tool_calls(tool_calls, my_executor)
-messages.extend(results)
-
-# Handle errors gracefully
-try:
- result = await my_tools[name](**args)
- return create_tool_result(call_id, serialize_tool_result(result))
-except Exception as e:
- return create_tool_error_result(call_id, str(e))
-```
+See the example repo's `tools.py` for the recommended pattern using these utilities.
---
## Message Utilities
-Helper functions for working with messages.
-
-```python
-from osmosis_ai.rollout import (
- parse_tool_calls,
- normalize_stop,
- get_message_content,
- get_message_role,
- is_assistant_message,
- is_tool_message,
- is_user_message,
- count_messages_by_role,
-)
-
-# Safely extract tool_calls from assistant message
-tool_calls = parse_tool_calls(assistant_message)
-if tool_calls:
- for tc in tool_calls:
- # Process tool call...
-
-# Normalize stop parameter (handles None, str, List[str])
-stop = normalize_stop(params.get("stop")) # Always List[str] or None
+The `osmosis_ai.rollout` module exports helpers for working with messages:
-# Get message content safely
-content = get_message_content(message)
-
-# Check message roles
-if is_assistant_message(message):
- pass
-elif is_tool_message(message):
- pass
-
-# Count messages by role
-counts = count_messages_by_role(messages)
-# {"user": 3, "assistant": 2, "tool": 1}
-```
+| Function | Description |
+|----------|-------------|
+| `parse_tool_calls(message)` | Safely extract `tool_calls` from an assistant message |
+| `normalize_stop(value)` | Normalize stop parameter to `List[str]` or `None` |
+| `get_message_content(message)` | Get message content safely |
+| `is_assistant_message(message)` | Check if message role is `assistant` |
+| `is_tool_message(message)` | Check if message role is `tool` |
+| `is_user_message(message)` | Check if message role is `user` |
+| `count_messages_by_role(messages)` | Count messages by role, returns `{"user": N, ...}` |
---
-## Using the Registry
-
-For applications with multiple agent types.
-
-```python
-from osmosis_ai.rollout import (
- register_agent_loop,
- get_agent_loop,
- list_agent_loops,
- create_app,
-)
-
-# Register agents at startup
-register_agent_loop(CalculatorAgent())
-register_agent_loop(SearchAgent())
-register_agent_loop(WeatherAgent())
-
-# List available agents
-print(list_agent_loops()) # ['calculator', 'search', 'weather']
-
-# Get specific agent
-agent = get_agent_loop("calculator")
-app = create_app(agent)
-```
-
-## Example Repository
-
-For a complete, runnable project with tools, rewards, and server setup, see: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
-
## See Also
- [Testing](./testing.md) -- unit tests and mock trainer
diff --git a/docs/remote-rollout/overview.md b/docs/remote-rollout/overview.md
index bc29cfa0..fee5318f 100644
--- a/docs/remote-rollout/overview.md
+++ b/docs/remote-rollout/overview.md
@@ -28,131 +28,61 @@ pip install osmosis-ai[full]
## Quick Start
+### Project Structure
+
+```
+rollout-server/
+├── server.py # Agent loop + FastAPI app
+├── tools.py # Tool definitions and execution
+├── rewards.py # Reward computation
+├── test_data.jsonl # Test dataset
+└── pyproject.toml # Dependencies: osmosis-ai[server]>=0.2.14
+```
+
+For the complete working project, see [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example).
+
### Step 1: Implement Your Agent Loop
-Create a class that inherits from `RolloutAgentLoop` and implements the required `get_tools` and `run` methods:
+Create `server.py` with a class that inherits from `RolloutAgentLoop`:
```python
-import json
-
-from osmosis_ai.rollout import (
- RolloutAgentLoop,
- RolloutContext,
- RolloutResult,
- RolloutRequest,
- OpenAIFunctionToolSchema,
- OpenAIFunctionSchema,
-)
+class CalculatorAgentLoop(RolloutAgentLoop):
+ name = "calculator"
-class MyAgentLoop(RolloutAgentLoop):
- name = "my_agent" # Required: unique identifier
-
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- """Return tools available for this rollout."""
- return [
- OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="search",
- description="Search for information",
- ),
- )
- ]
-
- async def execute_tool(self, name: str, args: dict) -> str:
- """Execute a tool and return the result as a string."""
- if name == "search":
- return f"Search results for '{args.get('query', '')}'"
- return f"Unknown tool: {name}"
+ def get_tools(self, request: RolloutRequest):
+ return CALCULATOR_TOOL_SCHEMAS
async def run(self, ctx: RolloutContext) -> RolloutResult:
- """Execute the agent loop."""
messages = list(ctx.request.messages)
-
- for _ in range(ctx.request.max_turns):
- # Call LLM through TrainGate
+ for _turn in range(ctx.request.max_turns):
result = await ctx.chat(messages, **ctx.request.completion_params)
messages.append(result.message)
-
- # Check if done
if not result.has_tool_calls:
break
+ tool_results = await execute_tools(result.tool_calls)
+ messages.extend(tool_results)
+ return ctx.complete(messages, reward=compute_reward(...))
- # Execute tools and add results
- for tool_call in result.tool_calls:
- tool_name = tool_call["function"]["name"]
- tool_args = json.loads(tool_call["function"]["arguments"])
-
- # Execute tool logic (replace with your implementation)
- tool_result = await self.execute_tool(tool_name, tool_args)
-
- messages.append({
- "role": "tool",
- "content": tool_result,
- "tool_call_id": tool_call["id"],
- })
- ctx.record_tool_call()
-
- return ctx.complete(messages)
+agent_loop = CalculatorAgentLoop()
+app = create_app(agent_loop)
```
-### Step 2: Run the Server
+See [Examples](./examples.md) for the complete `tools.py` and `rewards.py` files.
-#### Option 1: Using CLI (Recommended)
-
-Export an instance of your agent loop and use the CLI:
-
-```python
-# my_agent.py
-agent_loop = MyAgentLoop()
-```
+### Step 2: Run the Server
```bash
# Validate agent loop (checks tools, async run method, etc.)
-osmosis validate -m my_agent:agent_loop
+osmosis validate -m server:agent_loop
# Start server with validation (default port 9000)
-osmosis serve -m my_agent:agent_loop
-
-# Specify port
-osmosis serve -m my_agent:agent_loop -p 8080
+osmosis serve -m server:agent_loop
-# Skip validation (not recommended)
-osmosis serve -m my_agent:agent_loop --no-validate
-
-# Enable auto-reload for development
-osmosis serve -m my_agent:agent_loop --reload
+# Specify port and enable auto-reload
+osmosis serve -m server:agent_loop -p 8080 --reload
```
-#### Option 2: Using create_app()
-
-Create a FastAPI application manually:
-
-```python
-from osmosis_ai.rollout import create_app
-
-app = create_app(MyAgentLoop())
-```
-
-```bash
-uvicorn main:app --host 0.0.0.0 --port 9000
-```
-
-#### Option 3: Using serve_agent_loop()
-
-Start programmatically with validation:
-
-```python
-from osmosis_ai.rollout import serve_agent_loop
-
-# Validates and starts server
-serve_agent_loop(MyAgentLoop(), port=9000)
-
-# Skip validation (not recommended)
-serve_agent_loop(MyAgentLoop(), port=9000, validate=False)
-```
-
-Your server is now ready to receive rollout requests from TrainGate.
+For programmatic alternatives (`create_app()`, `serve_agent_loop()`), see [Agent Loop Guide](./agent-loop.md#server).
## Key Concepts
@@ -172,12 +102,7 @@ Provided to your `run()` method, containing:
- `tools`: List of tools returned by `get_tools()`
- `llm`: The `OsmosisLLMClient` for LLM calls
-Key methods:
-
-- `ctx.chat(messages, **kwargs)`: Call the LLM
-- `ctx.complete(messages)`: Return successful result
-- `ctx.error(message)`: Return error result
-- `ctx.record_tool_call()`: Track tool execution metrics
+Key methods: `ctx.chat()`, `ctx.complete()`, `ctx.error()`, `ctx.record_tool_call()`
### RolloutResult
@@ -188,26 +113,6 @@ The return value from your agent loop:
- `finish_reason`: Why the rollout ended
- `metrics`: Execution metrics (latency, token counts)
-### OpenAIFunctionToolSchema
-
-OpenAI-compatible tool definition format. Define your tools using:
-
-```python
-OpenAIFunctionToolSchema(
- type="function",
- function=OpenAIFunctionSchema(
- name="tool_name",
- description="What this tool does",
- parameters=OpenAIFunctionParametersSchema(
- properties={
- "arg1": OpenAIFunctionPropertySchema(type="string"),
- },
- required=["arg1"],
- ),
- ),
-)
-```
-
## Server Endpoints
The server created by `create_app()` exposes:
@@ -234,12 +139,7 @@ When installed with `osmosis-ai[server]`, settings can be loaded from environmen
- `OSMOSIS_ROLLOUT_CLIENT_*` - Client settings (timeout, retries, etc.)
- `OSMOSIS_ROLLOUT_SERVER_*` - Server settings (concurrency, TTL, etc.)
-Example:
-
-```bash
-export OSMOSIS_ROLLOUT_CLIENT_TIMEOUT_SECONDS=120
-export OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS=200
-```
+See [Configuration](../configuration.md) for the full list of environment variables.
## Example Repository
diff --git a/docs/remote-rollout/testing.md b/docs/remote-rollout/testing.md
index 429f4574..5aa13d28 100644
--- a/docs/remote-rollout/testing.md
+++ b/docs/remote-rollout/testing.md
@@ -62,90 +62,6 @@ def weather_tool_generator(message):
app = create_mock_trainer_app(tool_call_generator=weather_tool_generator)
```
-## Example Test Using FastAPI TestClient
-
-```python
-# test_my_agent.py
-
-import pytest
-from unittest.mock import AsyncMock, patch
-
-from fastapi.testclient import TestClient
-from osmosis_ai.rollout import create_app
-
-from my_agent import MyAgentLoop
-
-
-@pytest.fixture
-def app():
- """Create test application."""
- return create_app(MyAgentLoop())
-
-
-@pytest.fixture
-def client(app):
- """Create test client."""
- with TestClient(app) as client:
- yield client
-
-
-def test_health_endpoint(client):
- """Test health check."""
- response = client.get("/health")
- assert response.status_code == 200
- assert response.json()["status"] == "healthy"
-
-
-def test_init_returns_tools(client):
- """Test /v1/rollout/init returns tools."""
- response = client.post(
- "/v1/rollout/init",
- json={
- "rollout_id": "test-123",
- "server_url": "http://localhost:8080",
- "messages": [{"role": "user", "content": "Hello"}],
- "completion_params": {"temperature": 0.7},
- },
- )
- assert response.status_code == 202
- data = response.json()
- assert "tools" in data
-
-
-@pytest.mark.asyncio
-async def test_agent_run():
- """Test agent run directly."""
- from unittest.mock import MagicMock
- from osmosis_ai.rollout import RolloutContext, RolloutRequest, RolloutMetrics
-
- agent = MyAgentLoop()
-
- # Create mock context
- mock_llm = MagicMock()
- mock_llm.get_metrics.return_value = RolloutMetrics()
- mock_llm.chat_completions = AsyncMock(return_value=MagicMock(
- message={"role": "assistant", "content": "Hello!"},
- has_tool_calls=False,
- ))
-
- request = RolloutRequest(
- rollout_id="test",
- server_url="http://localhost",
- messages=[{"role": "user", "content": "Hi"}],
- completion_params={},
- )
-
- ctx = RolloutContext(
- request=request,
- tools=[],
- llm=mock_llm,
- )
-
- result = await agent.run(ctx)
-
- assert result.status == "COMPLETED"
-```
-
## Testing Utilities Reference
The `osmosis_ai.rollout.testing` module provides the following public utilities for testing agent loops without a real TrainGate server.
@@ -177,25 +93,11 @@ Thread-safe tracker that captures `/v1/rollout/completed` callbacks during tests
Patches `httpx.AsyncClient.post` so that any requests to `/v1/chat/completions` or `/v1/rollout/completed` are routed to the mock trainer `TestClient` instead of making real HTTP calls. All other requests pass through unchanged.
-```python
-@pytest.fixture
-def mock_trainer(monkeypatch):
- tracker = RolloutCompletionTracker()
- app = create_mock_trainer_app(tracker=tracker)
- client = TestClient(app)
- patch_httpx_for_mock_trainer(client, monkeypatch)
- return client, tracker
-```
-
-### `fake_token_ids(text)`
-
-Generates deterministic fake token IDs for testing. Returns a list of sequential integers, one per character in the input text (e.g. `fake_token_ids("hello")` returns `[0, 1, 2, 3, 4]`).
-
-### `fake_prompt_token_ids(messages)`
+### `fake_token_ids(text)` / `fake_prompt_token_ids(messages)`
-Generates deterministic fake prompt token IDs for testing. The token count scales with the number of messages to simulate realistic prompt growth. Returns `list(range(10 * max(1, len(messages))))`.
+Generate deterministic fake token IDs for testing. These produce deterministic output suitable for snapshot testing but do not correspond to any real tokenizer.
-> **Note:** These fake token ID functions produce deterministic output suitable for snapshot testing but do not correspond to any real tokenizer.
+For a complete test file example, see the [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) repository.
## See Also
diff --git a/docs/test-mode.md b/docs/test-mode.md
index dcb2a138..2af9f806 100644
--- a/docs/test-mode.md
+++ b/docs/test-mode.md
@@ -14,81 +14,37 @@ Test mode enables you to:
- Use 100+ LLM providers (OpenAI, Anthropic, Groq, Ollama, etc.)
- Run batch tests against datasets with detailed metrics
- Debug step-by-step with interactive mode
-- Track token usage, latency, and reward
## Quick Start
-### Testing with Remote Rollout Agent
-
-If you implement a `RolloutAgentLoop`, use `-m` to point at your agent module:
-
```bash
-# Basic batch test with OpenAI
-osmosis test -m my_agent:agent_loop -d data.jsonl --model gpt-5-mini
-
-# Use Anthropic Claude
-osmosis test -m my_agent:agent_loop -d data.jsonl --model anthropic/claude-sonnet-4-5
-
-# Interactive debugging
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
-
-# Start at specific row
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive --row 5
-```
-
-### Testing with Local Rollout MCP Tools
-
-If you use the **git-sync** workflow and provide MCP tools (`@mcp.tool()`) instead of a `RolloutAgentLoop`, use `--mcp` to point at your MCP directory:
+# Remote Rollout: batch test with OpenAI
+osmosis test -m server:agent_loop -d data.jsonl --model gpt-5-mini
-```bash
-# Install MCP support
-pip install osmosis-ai[mcp]
-
-# Test MCP tools against a dataset
+# Local Rollout: test MCP tools
osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-# Interactive debugging with MCP tools
-osmosis test --mcp ./mcp -d data.jsonl --interactive
-
-# All options work: limits, temperature, output, etc.
-osmosis test --mcp ./mcp -d data.jsonl --model gpt-5-mini \
- --limit 5 --temperature 0.7 -o results.json
+# Interactive debugging (start at specific row)
+osmosis test -m server:agent_loop -d data.jsonl --interactive --row 5
```
-The `--mcp` directory must contain a `main.py` that creates a `FastMCP` instance with registered tools. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) for the full folder structure and examples.
-
-> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use one or the other.
+The `--mcp` directory must contain a `main.py` with a `FastMCP` instance. See [Local Rollout MCP Tools](./local-rollout/mcp-tools.md). `--mcp` and `-m` are mutually exclusive.
### Programmatic Usage
```python
-from osmosis_ai.rollout.eval.common import DatasetReader
-from osmosis_ai.rollout.eval.common import ExternalLLMClient
+from osmosis_ai.rollout.eval.common import DatasetReader, ExternalLLMClient
from osmosis_ai.rollout.eval.test_mode import LocalTestRunner
-# Load dataset
reader = DatasetReader("./test_data.jsonl")
rows = reader.read(limit=10)
+client = ExternalLLMClient("gpt-5-mini")
-# Initialize LLM client
-client = ExternalLLMClient("gpt-5-mini") # or "anthropic/claude-sonnet-4-5"
-
-# Create runner
-runner = LocalTestRunner(
- agent_loop=MyAgentLoop(),
- llm_client=client,
-)
-
-# Run batch tests
+runner = LocalTestRunner(agent_loop=CalculatorAgentLoop(), llm_client=client)
async with client:
- results = await runner.run_batch(
- rows=rows,
- max_turns=10,
- completion_params={"temperature": 0.7},
- )
+ results = await runner.run_batch(rows=rows, max_turns=10)
print(f"Passed: {results.passed}/{results.total}")
-print(f"Total tokens: {results.total_tokens}")
```
---
@@ -115,8 +71,8 @@ osmosis test [OPTIONS]
| Option | Description |
|--------|-------------|
-| `-m, --module, --agent MODULE` | Module path to agent loop (format: `module:attribute`). Use for Remote Rollout users who implement `RolloutAgentLoop`. |
-| `--mcp DIR` | Path to MCP tools directory (must contain `main.py` with a `FastMCP` instance). Use for Local Rollout (git-sync) users who provide `@mcp.tool()` functions. Requires `pip install osmosis-ai[mcp]`. |
+| `-m, --module, --agent MODULE` | Module path to agent loop (format: `module:attribute`) |
+| `--mcp DIR` | Path to MCP tools directory. Requires `pip install osmosis-ai[mcp]`. |
### Model Options
@@ -160,36 +116,6 @@ Models can be specified in two formats:
See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for supported providers.
-### Examples
-
-```bash
-# Remote Rollout: test with GPT-5-mini (default)
-osmosis test -m my_agent:agent_loop -d data.jsonl
-
-# Local Rollout: test MCP tools
-osmosis test --mcp ./mcp -d data.jsonl --model openai/gpt-5-mini
-
-# Local Rollout: interactive debugging with MCP tools
-osmosis test --mcp ./mcp -d data.jsonl --interactive
-
-# Test with Claude
-osmosis test -m my_agent:agent_loop -d data.jsonl --model anthropic/claude-sonnet-4-5
-
-# Test with local Ollama
-osmosis test -m my_agent:agent_loop -d data.jsonl \
- --model ollama/llama3.1 \
- --base-url http://localhost:11434
-
-# Test subset of data
-osmosis test -m my_agent:agent_loop -d data.jsonl --limit 10 --offset 50
-
-# Save results to file
-osmosis test -m my_agent:agent_loop -d data.jsonl -o results.json
-
-# Interactive debugging starting at row 5
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive --row 5
-```
-
---
## Interactive Mode
@@ -208,185 +134,55 @@ Interactive mode allows step-by-step debugging of agent execution. After each LL
| `r` / `row N` | Jump to row N |
| `?` / `help` | Show help |
-### Example Session
-
-```
-osmosis-rollout-test v0.1.0 (Interactive Mode)
-
-Loading agent: my_agent:agent_loop
- Agent name: calculator
-Loading dataset: data.jsonl
- Total rows: 100
-Initializing provider: openai
- Model: openai/gpt-5-mini
-
-[Interactive Mode] Row 0/100
-System: You are a helpful calculator.
-User: What is 2 + 2?
-
-[Step 1] Waiting for LLM response...
-> n
-
-Assistant called tool: calculate(operation="add", a=2, b=2)
-Tool result: 4
-
-[Step 2] Waiting for LLM response...
-> m
-
-Messages:
- [0] system: You are a helpful calculator.
- [1] user: What is 2 + 2?
- [2] assistant: [tool_call: calculate]
- [3] tool: 4
-
-> c
-
-Continuing without stepping...
-Result: COMPLETED (reward=1.0)
-
-[Row 0 Complete] Next row? (n=next, q=quit, r N=jump to row N)
-> q
-```
-
---
## API Reference
### DatasetReader
-Read and validate test datasets.
-
```python
from osmosis_ai.rollout.eval.common import DatasetReader
reader = DatasetReader("./data.jsonl")
-
-# Get total row count
-total = len(reader)
-
-# Read all rows
-rows = reader.read()
-
-# Read with pagination
rows = reader.read(limit=10, offset=20)
-
-# Iterate rows (memory efficient)
-for row in reader.iter_rows():
- print(row["user_prompt"])
```
-**Constructor Parameters:**
-
-| Parameter | Type | Description |
-|-----------|------|-------------|
-| `path` | `str` | Path to dataset file |
-
-**Methods:**
-
| Method | Returns | Description |
|--------|---------|-------------|
| `read(limit, offset)` | `List[DatasetRow]` | Read rows with optional pagination |
| `iter_rows()` | `Iterator[DatasetRow]` | Memory-efficient iterator |
| `__len__()` | `int` | Total row count |
----
-
-### DatasetRow
-
-A single row from the dataset, represented as a `TypedDict`.
-
-```python
-class DatasetRow(TypedDict):
- ground_truth: str # Expected output
- user_prompt: str # User message
- system_prompt: str # System prompt
- # Additional columns from the dataset are included as extra dict keys
-```
-
-Access fields with dict syntax: `row["user_prompt"]`, `row["ground_truth"]`, etc.
-
----
+`DatasetRow` is a `TypedDict` with `ground_truth`, `user_prompt`, `system_prompt`, plus any extra columns from the dataset.
### ExternalLLMClient
-Call external LLM APIs via LiteLLM.
-
```python
from osmosis_ai.rollout.eval.common import ExternalLLMClient
-# OpenAI (simple format)
-client = ExternalLLMClient("gpt-5-mini")
-
-# Anthropic (LiteLLM format)
-client = ExternalLLMClient("anthropic/claude-sonnet-4-5")
-
-# With explicit API key
-client = ExternalLLMClient(
- model="gpt-5-mini",
- api_key="sk-...",
-)
-
-# Local Ollama
-client = ExternalLLMClient(
- model="ollama/llama3.1",
- api_base="http://localhost:11434",
-)
-
-# Use as async context manager
-async with client:
- result = await client.chat_completions(messages, tools=tools)
+client = ExternalLLMClient("gpt-5-mini") # OpenAI
+client = ExternalLLMClient("anthropic/claude-sonnet-4-5") # Anthropic
+client = ExternalLLMClient("ollama/llama3.1", api_base="http://localhost:11434") # Local
```
-**Constructor Parameters:**
-
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `model` | `str` | required | Model name (simple or LiteLLM format) |
| `api_key` | `Optional[str]` | `None` | API key (or use env var) |
| `api_base` | `Optional[str]` | `None` | Base URL for OpenAI-compatible APIs |
-**Methods:**
-
-#### `async chat_completions(messages, tools, **kwargs) -> CompletionsResult`
-
-Call the LLM with messages and tools.
-
-| Parameter | Type | Description |
-|-----------|------|-------------|
-| `messages` | `List[Dict]` | Conversation history |
-| `tools` | `List[Dict]` | Tool definitions (OpenAI format) |
-| `**kwargs` | - | Additional params (temperature, max_tokens, etc.) |
-
----
+Method: `async chat_completions(messages, tools, **kwargs) -> CompletionsResult`
### LocalTestRunner
-Execute batch tests.
-
```python
from osmosis_ai.rollout.eval.test_mode import LocalTestRunner
-runner = LocalTestRunner(
- agent_loop=MyAgentLoop(),
- llm_client=client,
- debug=True,
-)
-
-# Single row test
+runner = LocalTestRunner(agent_loop=agent, llm_client=client)
result = await runner.run_single(row, row_index=0)
-
-# Batch test
-batch_result = await runner.run_batch(
- rows=rows,
- max_turns=10,
- completion_params={"temperature": 0.7},
- on_progress=lambda current, total, result: print(f"[{current}/{total}]"),
- start_index=0,
-)
+batch_result = await runner.run_batch(rows=rows, max_turns=10)
```
-**Constructor Parameters:**
-
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `agent_loop` | `RolloutAgentLoop` | required | Agent to test |
@@ -394,97 +190,29 @@ batch_result = await runner.run_batch(
| `debug` | `bool` | `False` | Enable debug output |
| `debug_dir` | `Optional[str]` | `None` | Directory for debug logs |
-**Methods:**
-
-#### `async run_single(row, row_index, max_turns, completion_params) -> LocalTestRunResult`
-
-Test a single row.
-
-#### `async run_batch(rows, max_turns, completion_params, on_progress, start_index) -> LocalTestBatchResult`
-
-Test multiple rows.
-
----
-
-### LocalTestRunResult
-
-Result from a single test run.
-
-```python
-@dataclass
-class LocalTestRunResult:
- row_index: int # Dataset row index
- success: bool # Whether completed without error
- result: Optional[RolloutResult] # Agent result (if successful)
- error: Optional[str] # Error message (if failed)
- duration_ms: float # Execution time
- token_usage: Dict[str, int] # Token statistics
-```
-
-**Token Usage Fields:**
-
-| Field | Description |
-|-------|-------------|
-| `prompt_tokens` | Input tokens |
-| `completion_tokens` | Output tokens |
-| `total_tokens` | Total tokens |
-| `num_llm_calls` | Number of LLM calls |
+### Result Types
----
-
-### LocalTestBatchResult
+| Type | Key Fields |
+|------|------------|
+| `LocalTestRunResult` | `row_index`, `success`, `result: Optional[RolloutResult]`, `error`, `duration_ms`, `token_usage` |
+| `LocalTestBatchResult` | `results`, `total`, `passed`, `failed`, `total_duration_ms`, `total_tokens` |
-Aggregated results from batch testing.
-
-```python
-@dataclass
-class LocalTestBatchResult:
- results: List[LocalTestRunResult] # Individual results
- total: int # Total rows tested
- passed: int # Successful completions
- failed: int # Failed tests
- total_duration_ms: float # Total execution time
- total_tokens: int # Total tokens used
-```
-
----
+`token_usage` dict contains: `prompt_tokens`, `completion_tokens`, `total_tokens`, `num_llm_calls`.
### InteractiveRunner
-Step-by-step debugging runner.
-
```python
from osmosis_ai.rollout.eval.test_mode import InteractiveRunner
-runner = InteractiveRunner(
- agent_loop=MyAgentLoop(),
- llm_client=client,
- debug=True,
-)
-
-await runner.run_interactive_session(
- rows=rows,
- max_turns=10,
- completion_params={"temperature": 0.7},
- initial_row=5,
- row_offset=0,
-)
+runner = InteractiveRunner(agent_loop=agent, llm_client=client)
+await runner.run_interactive_session(rows=rows, max_turns=10, initial_row=5)
```
---
## Exceptions
-All local workflow exceptions are shared by local commands (`test` / `eval`).
-
-```python
-from osmosis_ai.rollout.eval.common import (
- DatasetValidationError,
- DatasetParseError,
- ProviderError,
- ToolValidationError,
-)
-```
+All local workflow exceptions are shared by `test` and `eval` commands.
| Exception | Description |
|-----------|-------------|
@@ -493,57 +221,14 @@ from osmosis_ai.rollout.eval.common import (
| `ProviderError` | LLM provider error (auth, rate limit, etc.) |
| `ToolValidationError` | Invalid tool schema |
-### Example Error Handling
-
-```python
-from osmosis_ai.rollout.eval.common import (
- DatasetReader,
- DatasetValidationError,
- DatasetParseError,
-)
-
-try:
- reader = DatasetReader("./data.jsonl")
- rows = reader.read()
-except DatasetParseError as e:
- print(f"Failed to parse file: {e}")
-except DatasetValidationError as e:
- print(f"Invalid dataset: {e}")
-```
-
---
## Output Format
-When using `--output`, results are saved as JSON:
-
-```json
-{
- "summary": {
- "total": 100,
- "passed": 95,
- "failed": 5,
- "total_duration_ms": 45230,
- "total_tokens": 125000
- },
- "results": [
- {
- "row_index": 0,
- "success": true,
- "error": null,
- "duration_ms": 450,
- "token_usage": {
- "prompt_tokens": 150,
- "completion_tokens": 50,
- "total_tokens": 200,
- "num_llm_calls": 2
- },
- "reward": 1.0,
- "finish_reason": "stop"
- }
- ]
-}
-```
+When using `--output`, results are saved as JSON with this structure:
+
+- `summary` -- total, passed, failed, total_duration_ms, total_tokens
+- `results[]` -- per-row: row_index, success, error, duration_ms, token_usage, reward, finish_reason
---
@@ -562,20 +247,6 @@ See [LiteLLM Environment Variables](https://docs.litellm.ai/docs/providers) for
---
-## Best Practices
-
-1. **Start with a small dataset**: Use `--limit 5` to validate your setup before running full tests.
-
-2. **Use interactive mode for debugging**: When a test fails, use `--interactive --row N` to step through execution.
-
-3. **Save results for analysis**: Use `-o results.json` to save detailed metrics for later analysis.
-
-4. **Set appropriate timeouts**: Complex agent loops may need longer `--max-turns` values.
-
-5. **Monitor token usage**: Track `total_tokens` to estimate costs before running large batches.
-
----
-
## See Also
- [Eval Mode](./eval-mode.md) -- Evaluate agents with eval functions and pass@k
diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md
index 88095d4a..1ebd841d 100644
--- a/docs/troubleshooting.md
+++ b/docs/troubleshooting.md
@@ -295,7 +295,7 @@ Common validation error codes:
Run validation before serving to catch these early:
```bash
-osmosis validate -m my_agent:agent_loop
+osmosis validate -m server:agent_loop
```
### ServeError
@@ -381,11 +381,11 @@ as a fallback. If all fail:
### Using `--verbose` flag
-The `osmosis serve` command accepts `-v` / `--verbose` to increase output
-verbosity, which prints detailed validation warnings and server configuration.
+The `osmosis validate` command accepts `-v` / `--verbose` to show detailed
+validation output including warnings.
```bash
-osmosis serve -m my_agent:agent_loop --verbose
+osmosis validate -m server:agent_loop --verbose
```
### Using `--debug` flag
@@ -395,8 +395,8 @@ logging, which prints detailed information about each step including provider
requests and responses.
```bash
-osmosis test -m my_agent:agent_loop -d data.jsonl --model openai/gpt-5-mini --debug
-osmosis eval -m my_agent:agent_loop -d data.jsonl --eval-fn rewards:fn --model openai/gpt-5-mini --debug
+osmosis test -m server:agent_loop -d data.jsonl --model openai/gpt-5-mini --debug
+osmosis eval -m server:agent_loop -d data.jsonl --eval-fn rewards:fn --model openai/gpt-5-mini --debug
```
### Interactive mode in test mode
@@ -405,19 +405,19 @@ Use `--interactive` with `osmosis test` to step through each row one at a time.
This is useful for debugging agent logic and inspecting intermediate messages:
```bash
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive
+osmosis test -m server:agent_loop -d data.jsonl --interactive
# Start at a specific row
-osmosis test -m my_agent:agent_loop -d data.jsonl --interactive --row 5
+osmosis test -m server:agent_loop -d data.jsonl --interactive --row 5
```
### Debug directory for rollout server
-When serving an agent, use `--debug-dir` to write detailed execution traces
+When serving an agent, use `--log` to write detailed execution traces
for each rollout to JSONL files:
```bash
-osmosis serve -m my_agent:agent_loop --debug-dir ./debug-logs
+osmosis serve -m server:agent_loop --log ./debug-logs
```
Each rollout will produce a file at `{debug_dir}/{timestamp}/{rollout_id}.jsonl`.
@@ -435,7 +435,7 @@ logging.basicConfig(level=logging.DEBUG)
Or set the uvicorn log level when serving:
```bash
-osmosis serve -m my_agent:agent_loop --log-level debug
+osmosis serve -m server:agent_loop --log-level debug
```
### Checking credentials
From b22c8f372b1d516fecbf76c9b78c4b0fdc754e39 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 16:10:40 -0800
Subject: [PATCH 43/60] simplify codeowners
---
.github/CODEOWNERS | 14 +-------------
1 file changed, 1 insertion(+), 13 deletions(-)
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index cc3ad7c6..c26bc239 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,13 +1 @@
-# Default: core maintainers review everything
-* @BaiqingL @JoyboyBrian @JakeTrock @allenlinsh @mathewjhan @pandyamarut @artem-osmosis
-
-# Rollout SDK (core training integration)
-/osmosis_ai/rollout/ @BaiqingL @JoyboyBrian @JakeTrock @allenlinsh @mathewjhan @pandyamarut @artem-osmosis
-
-# CI/CD and build configuration
-/.github/ @BaiqingL @JoyboyBrian @JakeTrock
-/pyproject.toml @BaiqingL @JoyboyBrian @JakeTrock
-
-# Documentation
-/README.md @BaiqingL @JoyboyBrian @JakeTrock
-/CONTRIBUTING.md @BaiqingL @JoyboyBrian @JakeTrock
+* @BaiqingL @JoyboyBrian
From 93c594eb10da3f3aeb04f5cbbfcf19b4c1eac38f Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 16:12:47 -0800
Subject: [PATCH 44/60] error handling
---
osmosis_ai/rollout/eval/common/dataset.py | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/osmosis_ai/rollout/eval/common/dataset.py b/osmosis_ai/rollout/eval/common/dataset.py
index 51edd002..25d3884b 100644
--- a/osmosis_ai/rollout/eval/common/dataset.py
+++ b/osmosis_ai/rollout/eval/common/dataset.py
@@ -130,8 +130,8 @@ def _count_jsonl_rows(self) -> int:
for line in f:
if line.strip():
count += 1
- except OSError:
- pass
+ except OSError as e:
+ logger.warning("Failed to count JSONL rows in %s: %s", self.file_path, e)
return count
def _parse_parquet(self) -> list[dict[str, Any]]:
@@ -187,8 +187,8 @@ def _count_csv_rows(self) -> int:
reader = csv.DictReader(f)
for _ in reader:
count += 1
- except (OSError, csv.Error):
- pass
+ except (OSError, csv.Error) as e:
+ logger.warning("Failed to count CSV rows in %s: %s", self.file_path, e)
return count
def _validate_row(self, row: Any, row_index: int) -> DatasetRow:
From 9010320b3547636ac73141a32152ff08531cc9a1 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 16:31:24 -0800
Subject: [PATCH 45/60] Refactor and enhance documentation for clarity and
structure
- Removed redundant sections and streamlined content across various documentation files, including the README and local rollout guides.
- Updated the README to better outline the Local and Remote Rollout modes, including quick start links and relevant documentation.
- Clarified the usage of reward functions and rubric evaluators, ensuring consistency in examples and explanations.
- Improved error handling documentation and removed outdated references to enhance user experience.
---
README.md | 49 +----
docs/README.md | 56 ++----
docs/cli.md | 2 +-
docs/configuration.md | 2 +-
docs/eval-mode.md | 129 +-----------
docs/local-rollout/mcp-tools.md | 59 +-----
docs/local-rollout/overview.md | 34 +---
docs/local-rollout/reward-functions.md | 21 +-
docs/local-rollout/reward-rubrics.md | 88 ++++-----
docs/remote-rollout/agent-loop.md | 261 ++-----------------------
docs/remote-rollout/architecture.md | 149 +-------------
docs/remote-rollout/examples.md | 56 +-----
docs/remote-rollout/overview.md | 121 +-----------
docs/remote-rollout/testing.md | 105 ----------
docs/rewards-api.md | 222 ---------------------
docs/test-mode.md | 121 +-----------
docs/troubleshooting.md | 232 ++--------------------
17 files changed, 120 insertions(+), 1587 deletions(-)
delete mode 100644 docs/remote-rollout/testing.md
delete mode 100644 docs/rewards-api.md
diff --git a/README.md b/README.md
index 7b6150b4..e3d7ed86 100644
--- a/README.md
+++ b/README.md
@@ -67,58 +67,19 @@ uv add osmosis-ai[full] # All features
## Local Rollout
-Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools via a GitHub-synced repo. Best for standard tool-use agents, fast iteration, and zero infrastructure.
+Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools via a GitHub-synced repo.
-Get started with the example repo: **[osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)**
-
-For details, see the [Local Rollout docs](docs/local-rollout/overview.md).
+Get started: **[osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)** | [Docs](docs/local-rollout/overview.md)
## Remote Rollout
-You implement and host a `RolloutAgentLoop` server. Full control over agent behavior, custom orchestration, and persistent environments.
-
-Get started with the example repo: **[osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)**
+You implement and host a `RolloutAgentLoop` server. Full control over agent behavior.
-For details, see the [Remote Rollout docs](docs/remote-rollout/overview.md).
+Get started: **[osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)** | [Docs](docs/remote-rollout/overview.md)
## Testing & Evaluation
-Both modes share the same `osmosis test` and `osmosis eval` CLI tools. See the example repos for usage:
-
-- **Local Rollout**: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
-- **Remote Rollout**: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
-
-For CLI details, see [Test Mode](docs/test-mode.md), [Eval Mode](docs/eval-mode.md), and the [CLI Reference](docs/cli.md).
-
-## Documentation
-
-| Section | Topics |
-|---------|--------|
-| **Local Rollout** | |
-| [Overview](docs/local-rollout/overview.md) | When to choose, repo structure, setup |
-| [Reward Functions](docs/local-rollout/reward-functions.md) | `@osmosis_reward` decorator |
-| [Reward Rubrics](docs/local-rollout/reward-rubrics.md) | `@osmosis_rubric`, `evaluate_rubric` |
-| [MCP Tools](docs/local-rollout/mcp-tools.md) | `@mcp.tool()` definitions |
-| **Remote Rollout** | |
-| [Overview](docs/remote-rollout/overview.md) | Quick start, key concepts |
-| [Architecture](docs/remote-rollout/architecture.md) | Protocol design and lifecycle |
-| [Agent Loop Guide](docs/remote-rollout/agent-loop.md) | API reference |
-| [Examples](docs/remote-rollout/examples.md) | Agent implementations |
-| [Testing](docs/remote-rollout/testing.md) | Unit tests and mock trainer |
-| [Deployment](docs/remote-rollout/deployment.md) | Docker, production config |
-| **Shared** | |
-| [Dataset Format](docs/datasets.md) | Parquet, JSONL, CSV formats |
-| [Test Mode](docs/test-mode.md) | `osmosis test` |
-| [Eval Mode](docs/eval-mode.md) | `osmosis eval` with pass@k |
-| **Reference** | |
-| [Rewards API](docs/rewards-api.md) | Decorator and function signatures |
-| [CLI Reference](docs/cli.md) | All `osmosis` commands |
-
-Full documentation index: [docs/README.md](docs/README.md)
-
-## Examples
-
-The [`examples/`](examples/) directory contains standalone SDK API examples (reward functions, rubric evaluation). See [`examples/README.md`](examples/README.md) for details.
+Both modes share the same CLI tools: [Test Mode](docs/test-mode.md) | [Eval Mode](docs/eval-mode.md) | [CLI Reference](docs/cli.md)
## Contributing
diff --git a/docs/README.md b/docs/README.md
index dfbb3acf..0f97b4fb 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,52 +1,24 @@
# Osmosis SDK Documentation
-Python SDK for Osmosis AI training workflows. The SDK supports two parallel training modes -- **Local Rollout** and **Remote Rollout** -- along with shared tools for testing, evaluation, and reward scoring.
-
-## Getting Started
-
-- [Installation & Quick Start](../README.md)
-- [Contributing](../CONTRIBUTING.md)
-
## Local Rollout
-Git-sync mode where Osmosis manages the agent loop. You provide reward functions, rubrics, and MCP tools.
-
-- [Overview](./local-rollout/overview.md) -- what is local rollout, when to choose it, repository structure
-- [Reward Functions](./local-rollout/reward-functions.md) -- `@osmosis_reward` decorator usage and examples
-- [Reward Rubrics](./local-rollout/reward-rubrics.md) -- `@osmosis_rubric` and `evaluate_rubric` usage
-- [MCP Tools](./local-rollout/mcp-tools.md) -- defining `@mcp.tool()` functions for the agent
+- [Overview](./local-rollout/overview.md)
+- [Reward Functions](./local-rollout/reward-functions.md)
+- [Reward Rubrics](./local-rollout/reward-rubrics.md)
+- [MCP Tools](./local-rollout/mcp-tools.md)
## Remote Rollout
-Self-hosted mode where you implement and run the agent loop as a server.
-
-- [Overview](./remote-rollout/overview.md) -- quick start and key concepts
-- [Architecture](./remote-rollout/architecture.md) -- protocol design and agent lifecycle
-- [Agent Loop Guide](./remote-rollout/agent-loop.md) -- API reference for classes, schemas, types
-- [Examples](./remote-rollout/examples.md) -- agent implementations and utilities
-- [Testing](./remote-rollout/testing.md) -- unit tests and mock trainer
-
-## Shared Concepts
-
-Used by both Local Rollout and Remote Rollout modes.
-
-- [Dataset Format](./datasets.md) -- supported formats and required columns
-- [Test Mode](./test-mode.md) -- test agents with cloud LLMs (`osmosis test`)
-- [Eval Mode](./eval-mode.md) -- evaluate agents with eval functions and pass@k (`osmosis eval`)
+- [Overview](./remote-rollout/overview.md)
+- [Architecture](./remote-rollout/architecture.md)
+- [Agent Loop Guide](./remote-rollout/agent-loop.md)
+- [Examples](./remote-rollout/examples.md)
## Reference
-- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric` API details
-- [CLI Reference](./cli.md) -- all `osmosis` commands
-- [Configuration](./configuration.md) -- environment variables, settings classes, and programmatic configuration
-- [Troubleshooting](./troubleshooting.md) -- common errors, debug tips, and resolutions
-
-## Contributing
-
-We welcome contributions! See the [Contributing Guide](../CONTRIBUTING.md) for development setup, coding standards, and pull request guidelines.
-
-## Other
-
-- [Example Code](../examples/) -- reward functions, rubric configs, sample data
-- [Security Policy](../SECURITY.md)
-- [License](../LICENSE)
+- [Dataset Format](./datasets.md)
+- [Test Mode](./test-mode.md)
+- [Eval Mode](./eval-mode.md)
+- [CLI Reference](./cli.md)
+- [Configuration](./configuration.md)
+- [Troubleshooting](./troubleshooting.md)
diff --git a/docs/cli.md b/docs/cli.md
index d2380954..d17a8372 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -177,5 +177,5 @@ osmosis eval-rubric --rubric support_followup --data examples/sample_data.jsonl
- [Test Mode](./test-mode.md) -- full `osmosis test` documentation
- [Eval Mode](./eval-mode.md) -- full `osmosis eval` documentation
-- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
+- [Reward Rubrics](./local-rollout/reward-rubrics.md) -- `@osmosis_rubric`, `evaluate_rubric`
- [Dataset Format](./datasets.md) -- supported formats and required columns
diff --git a/docs/configuration.md b/docs/configuration.md
index df8ba6d7..6c6f2d07 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -42,7 +42,7 @@ Server-side configuration for the RolloutServer process.
## Programmatic Configuration
-Configuration is organized into three Pydantic Settings classes that automatically read from environment variables and `.env` files:
+Configuration is organized into three Pydantic Settings classes. When the `pydantic-settings` package is installed (included in the `server` extra), these classes automatically read from environment variables and `.env` files. Without `pydantic-settings`, they fall back to plain Pydantic models that only accept values passed programmatically:
| Class | Env Prefix | Description |
|-------|-----------|-------------|
diff --git a/docs/eval-mode.md b/docs/eval-mode.md
index f3cf335e..1a42915b 100644
--- a/docs/eval-mode.md
+++ b/docs/eval-mode.md
@@ -2,23 +2,6 @@
Evaluate your trained models by running agent implementations against datasets with custom eval functions, statistical analysis, and pass@k metrics. Works with both **Local Rollout** (MCP tools) and **Remote Rollout** (RolloutAgentLoop) agents.
-## Overview
-
-Eval mode is designed for **evaluating trained models**. Connect to any OpenAI-compatible model serving endpoint (such as osmosis-serving, vLLM, or SGLang) and measure agent performance with custom eval functions.
-
-Key capabilities:
-
-- **Benchmark trained models** against eval functions by connecting to serving endpoints
-- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp`
-- Run multiple trials per row for pass@k analysis
-- Use existing `@osmosis_reward` functions or full-context eval functions
-- Get statistical summaries (mean, std, min, max) per eval function
-- Compare model quality with `--baseline-model`
-- **Concurrent execution** with `--batch-size` for faster benchmarks
-
-> **Note:** Command split (development-stage breaking change):
-> `osmosis eval-rubric` is for hosted rubric evaluation, while `osmosis eval` is for agent eval mode documented here.
-
## Quick Start
```bash
@@ -58,13 +41,7 @@ Eval mode works with any OpenAI-compatible serving endpoint via `--base-url`:
| SGLang | `http://localhost:30000/v1` |
| Ollama | `http://localhost:11434/v1` |
-The `--model` parameter should match the model name as registered in the serving endpoint. You can also use LiteLLM providers (e.g., `openai/gpt-5-mini`, `anthropic/claude-sonnet-4-5`) as either the primary or baseline model.
-
----
-
-## Dataset Format
-
-See [Dataset Format](./datasets.md) for supported formats and required columns.
+The `--model` parameter should match the model name as registered in the serving endpoint. You can also use LiteLLM providers (e.g., `openai/gpt-5-mini`) as either the primary or baseline model.
---
@@ -74,21 +51,16 @@ Eval functions score agent outputs. Two signatures are supported, auto-detected
### Simple Mode (compatible with `@osmosis_reward`)
-Use when you only need the final assistant response:
-
```python
def exact_match(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
```
- First parameter must be named `solution_str`
-- `extra_info` contains the full row dict (all columns)
- Compatible with functions decorated with `@osmosis_reward`
### Full Mode
-Use when you need the complete conversation history:
-
```python
def conversation_quality(messages: list, ground_truth: str, metadata: dict, **kwargs) -> float:
assistant_messages = [m for m in messages if m["role"] == "assistant"]
@@ -104,14 +76,7 @@ def conversation_quality(messages: list, ground_truth: str, metadata: dict, **kw
When `--n` is greater than 1, eval mode runs each dataset row multiple times and computes pass@k metrics.
-Formula: `pass@k = 1 - C(n-c, k) / C(n, k)`
-
-Where:
-- `n` = total runs per row
-- `c` = number of passing runs (score >= `--pass-threshold`)
-- `k` = sample size
-
-pass@k is computed per row, then averaged across all rows.
+Formula: `pass@k = 1 - C(n-c, k) / C(n, k)` where `n` = total runs, `c` = passing runs (score >= `--pass-threshold`), `k` = sample size.
---
@@ -148,7 +113,7 @@ osmosis eval [OPTIONS]
| Option | Description |
|--------|-------------|
-| `--baseline-model MODEL` | Baseline model for comparison (reports per-model summary statistics) |
+| `--baseline-model MODEL` | Baseline model for comparison |
| `--baseline-base-url URL` | Base URL for the baseline model's endpoint |
| `--baseline-api-key KEY` | API key for the baseline model provider |
@@ -175,97 +140,13 @@ osmosis eval [OPTIONS]
---
-## API Reference
-
-### EvalRunner
-
-Orchestrates benchmark execution with eval function scoring.
-
-```python
-from osmosis_ai.rollout.eval.evaluation import EvalRunner
-
-runner = EvalRunner(
- agent_loop=agent,
- llm_client=client,
- eval_fns=eval_fns,
- baseline_llm_client=baseline_client, # optional
-)
-
-result = await runner.run_eval(
- rows=rows, n_runs=5, max_turns=10,
- pass_threshold=0.5, batch_size=5,
- completion_params={"temperature": 0.7},
-)
-```
-
-**Constructor Parameters:**
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `agent_loop` | `RolloutAgentLoop` | required | Agent to benchmark |
-| `llm_client` | `ExternalLLMClient` | required | LLM client |
-| `eval_fns` | `List[EvalFnWrapper]` | required | Eval functions to apply |
-| `debug` | `bool` | `False` | Enable debug logging |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logs |
-| `llm_client_factory` | `Optional[Callable]` | `None` | Factory for concurrent LLM client instances |
-| `baseline_llm_client` | `Optional[ExternalLLMClient]` | `None` | Baseline LLM client for comparison mode |
-| `baseline_llm_client_factory` | `Optional[Callable]` | `None` | Factory for concurrent baseline LLM client instances |
-
-**Methods:**
-
-- `async run_single(row, row_index, run_index, max_turns, completion_params) -> EvalRunResult`
-- `async run_eval(rows, n_runs, max_turns, completion_params, pass_threshold, on_progress, start_index, batch_size) -> EvalResult`
-
-### Result Types
-
-| Type | Key Fields |
-|------|------------|
-| `EvalRunResult` | `run_index`, `success`, `scores: Dict[str, float]`, `duration_ms`, `tokens`, `error`, `model_tag` |
-| `EvalRowResult` | `row_index`, `runs: List[EvalRunResult]` |
-| `EvalResult` | `rows`, `eval_summaries`, `total_rows`, `total_runs`, `total_tokens`, `n_runs`, `pass_threshold`, `model_summaries` |
-| `EvalEvalSummary` | `mean`, `std`, `min`, `max`, `pass_at_k: Dict[int, float]` |
-| `EvalModelSummary` | `model`, `model_tag`, `eval_summaries`, `total_runs`, `total_tokens`, `total_duration_ms` |
-
-### load_eval_fns
-
-```python
-from osmosis_ai.rollout.eval.evaluation import load_eval_fns
-
-eval_fns = load_eval_fns(["rewards:exact_match", "rewards:partial_match"])
-```
-
-Each path must be in `module:function` format. Raises `EvalFnError` if loading fails.
-
----
-
-## Output Format
-
-When using `--output`, results are saved as JSON with this structure:
-
-- `config` -- model, n_runs, pass_threshold, eval_fns, baseline_model
-- `summary` -- total_rows, total_runs, per-eval stats (mean/std/min/max/pass_at_k), total_tokens
-- `rows[]` -- per-row results with individual run scores, duration, tokens, and `model_tag`
-- `model_summaries[]` -- per-model stats (only when `--baseline-model` is used)
-
----
-
## Exceptions
-All exceptions are shared with [Test Mode](./test-mode.md#exceptions).
-
| Exception | Description |
|-----------|-------------|
| `EvalFnError` | Eval function loading, signature detection, or execution error |
-| `DatasetValidationError` | Dataset missing required columns or invalid values |
-| `DatasetParseError` | Failed to parse dataset file |
-| `ProviderError` | LLM provider error (auth, rate limit, etc.) |
-| `ToolValidationError` | Invalid tool schema |
-
----
-
-## Environment Variables
-See [Test Mode - Environment Variables](./test-mode.md#environment-variables) for API key configuration.
+All other exceptions are shared with [Test Mode](./test-mode.md#exceptions).
---
@@ -273,6 +154,4 @@ See [Test Mode - Environment Variables](./test-mode.md#environment-variables) fo
- [Test Mode](./test-mode.md) -- Test agent logic with external LLMs
- [Dataset Format](./datasets.md) -- Supported formats and required columns
-- [Remote Rollout Examples](./remote-rollout/examples.md) -- Working code examples
-- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) -- MCP tool definition
- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) -- Supported LLM providers
diff --git a/docs/local-rollout/mcp-tools.md b/docs/local-rollout/mcp-tools.md
index 784ea385..4eebfead 100644
--- a/docs/local-rollout/mcp-tools.md
+++ b/docs/local-rollout/mcp-tools.md
@@ -21,7 +21,7 @@ def multiply(first_val: float, second_val: float) -> float:
return round(first_val * second_val, 4)
```
-The tool name, description, and parameter schemas are derived from the function name, docstring, and type annotations. Keep docstrings clear and descriptive -- the LLM uses them to decide when and how to call your tools.
+Keep docstrings clear and descriptive -- the LLM uses them to decide when and how to call your tools.
## Folder Structure
@@ -31,54 +31,15 @@ The `mcp/` directory must contain a `main.py` entry point that creates a `FastMC
mcp/
├── main.py # Entry point
├── server/
-│ ├── __init__.py
│ └── mcp_server.py # Creates the FastMCP instance
└── tools/
├── __init__.py # Imports all tool modules
- ├── math.py # @mcp.tool() functions
- ├── api_helpers.py
- └── ml_utils.py
+ └── math.py # @mcp.tool() functions
```
-### main.py
-
-```python
-# mcp/main.py
-from server import mcp
-from tools import * # Triggers @mcp.tool() registrations
-
-if __name__ == "__main__":
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--host', default='0.0.0.0')
- parser.add_argument('--port', type=int, default=8080)
- args = parser.parse_args()
- mcp.run(transport="http", host=args.host, port=args.port)
-```
-
-### server/mcp_server.py
-
-```python
-# mcp/server/mcp_server.py
-from fastmcp import FastMCP
-
-mcp = FastMCP("OsmosisTools")
-```
-
-### tools/__init__.py
-
-```python
-# mcp/tools/__init__.py
-from .math import *
-from .api_helpers import *
-```
-
-## Simpler Structure
-
For small projects, you can put everything in a single `main.py`:
```python
-# mcp/main.py
from fastmcp import FastMCP
mcp = FastMCP("my_tools")
@@ -87,26 +48,17 @@ mcp = FastMCP("my_tools")
def add(a: float, b: float) -> str:
"""Add two numbers."""
return str(a + b)
-
-@mcp.tool()
-def lookup(key: str) -> str:
- """Look up a value by key."""
- data = {"pi": "3.14159", "e": "2.71828"}
- return data.get(key, "not found")
```
## Testing MCP Tools Locally
-Use `osmosis test` with the `--mcp` flag to test your tools against a dataset:
-
```bash
-# Install MCP support
pip install "osmosis-ai[mcp]"
# Batch test
osmosis test --mcp ./mcp -d test_data.jsonl --model openai/gpt-5-mini
-# Interactive debugging -- step through each LLM call
+# Interactive debugging
osmosis test --mcp ./mcp -d test_data.jsonl --interactive
# Evaluate with a reward function
@@ -123,17 +75,14 @@ When you pass `--mcp ./mcp`, the SDK:
2. Discovers registered tools and converts them to OpenAI function-calling schemas
3. Runs a built-in agent loop that calls the LLM, executes tool calls against your MCP functions, and repeats until the LLM stops calling tools or `--max-turns` is reached
-This means you can iterate on tools locally, then push to GitHub for Osmosis to sync -- no `RolloutAgentLoop` code needed.
-
> **Note:** `--mcp` and `-m/--module` are mutually exclusive. Use `--mcp` for Local Rollout projects; use `-m` for Remote Rollout projects that implement `RolloutAgentLoop`.
## Example Repository
-See the complete working example with MCP tools, reward functions, and rubrics: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
+See the complete working example: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
## See Also
- [Test Mode](../test-mode.md) -- full documentation for `osmosis test`
- [Eval Mode](../eval-mode.md) -- full documentation for `osmosis eval`
- [Local Rollout Overview](./overview.md) -- repository structure and setup
-- [Dataset Format](../datasets.md) -- dataset format for testing and evaluation
diff --git a/docs/local-rollout/overview.md b/docs/local-rollout/overview.md
index b66ee5a4..13c9a324 100644
--- a/docs/local-rollout/overview.md
+++ b/docs/local-rollout/overview.md
@@ -1,24 +1,6 @@
# Local Rollout
-Local Rollout is one of two training modes supported by the Osmosis platform. In this mode, Osmosis manages the entire agent loop -- you provide **reward functions**, **rubric evaluators**, and optionally **MCP tools** via a git-sync repository. The training infrastructure handles LLM inference, tool execution, and trajectory collection automatically.
-
-## When to Choose Local Rollout
-
-| Consideration | Local Rollout | Remote Rollout |
-|---------------|--------------|----------------|
-| **Agent logic** | Managed by Osmosis | You implement `RolloutAgentLoop` |
-| **Tool execution** | MCP tools synced from GitHub | Custom tool code in your server |
-| **Infrastructure** | None -- Osmosis handles everything | You host a RolloutServer |
-| **Flexibility** | Standard agent loop, configurable tools | Full control over agent behavior |
-| **Best for** | Tool-use tasks with standard agent patterns | Custom agent architectures, complex orchestration |
-
-Choose **Local Rollout** when:
-- Your task follows a standard tool-use agent pattern (LLM calls tools in a loop)
-- You want zero infrastructure to manage
-- Your tools can be expressed as MCP `@mcp.tool()` functions
-- You want fast iteration via git push
-
-Choose **Remote Rollout** when you need full control over the agent loop, custom orchestration logic, or tools that require a persistent server environment. See the [Remote Rollout overview](../remote-rollout/overview.md).
+Local Rollout is one of two training modes supported by the Osmosis platform. In this mode, Osmosis manages the entire agent loop -- you provide **reward functions**, **rubric evaluators**, and optionally **MCP tools** via a git-sync repository.
## How It Works
@@ -36,28 +18,20 @@ Choose **Remote Rollout** when you need full control over the agent loop, custom
your-repo/
├── mcp/ # MCP tools (agent's callable functions)
│ ├── main.py # Entry point -- creates FastMCP instance
-│ ├── server/ # Server setup
-│ │ └── mcp_server.py
│ └── tools/ # Tool implementations
-│ ├── __init__.py
-│ └── math.py
├── reward_fn/ # Reward functions (deterministic scoring)
│ └── compute_reward.py
├── reward_rubric/ # Rubric evaluators (LLM-as-judge scoring)
-│ ├── reward_rubric_openai.py
-│ └── reward_rubric_anthropic.py
+│ └── reward_rubric_openai.py
├── test_data.jsonl # Sample dataset for local testing
└── pyproject.toml
```
-All three directories (`mcp/`, `reward_fn/`, `reward_rubric/`) are optional -- include only what your training run needs.
+All three directories are optional -- include only what your training run needs.
## Local Testing
-Before pushing to GitHub, test your setup locally:
-
```bash
-# Install MCP support
pip install "osmosis-ai[mcp]"
# Test MCP tools against a dataset
@@ -72,8 +46,6 @@ osmosis eval --mcp ./mcp -d test_data.jsonl \
--model openai/gpt-5-mini
```
-See [Test Mode](../test-mode.md) and [Eval Mode](../eval-mode.md) for full documentation.
-
## Example Repository
See the complete working example: [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example)
diff --git a/docs/local-rollout/reward-functions.md b/docs/local-rollout/reward-functions.md
index f000d71e..3bbcf9f1 100644
--- a/docs/local-rollout/reward-functions.md
+++ b/docs/local-rollout/reward-functions.md
@@ -1,6 +1,6 @@
# Reward Functions
-Reward functions provide deterministic, code-based scoring for LLM outputs during training. They are used in Local Rollout mode to drive reinforcement learning -- the training loop calls your reward function after each rollout to compute a scalar score.
+Reward functions provide deterministic, code-based scoring for LLM outputs during training.
## @osmosis_reward
@@ -22,10 +22,6 @@ def your_function(solution_str: str, ground_truth: str, extra_info: dict = None,
- **`extra_info: dict = None`** -- Optional dictionary for additional configuration.
- **`**kwargs`** -- Required for platform compatibility.
-### Return Value
-
-- **`-> float`** -- Must return a float value representing the reward score.
-
The decorator raises a `TypeError` if the function doesn't match this exact signature or doesn't return a float.
## Example
@@ -37,24 +33,11 @@ def exact_match(solution_str: str, ground_truth: str, extra_info: dict = None, *
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
```
-For more examples (numeric tolerance, solution extraction, partial credit), see:
+For more examples, see:
- [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) -- `rewards.py`
- [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) -- `reward_fn/compute_reward.py`
-## File Placement
-
-In a Local Rollout repository, place reward functions in the `reward_fn/` directory:
-
-```
-reward_fn/
-└── compute_reward.py # Contains @osmosis_reward functions
-```
-
-Osmosis discovers and syncs all `@osmosis_reward`-decorated functions from this directory.
-
## See Also
- [Reward Rubrics](./reward-rubrics.md) -- LLM-as-judge scoring with `@osmosis_rubric`
-- [Rewards API Reference](../rewards-api.md) -- full API reference for decorators and `evaluate_rubric`
- [Local Rollout Overview](./overview.md) -- repository structure and setup
-- [Example Code](../../examples/) -- reward functions, rubric configs, sample data
diff --git a/docs/local-rollout/reward-rubrics.md b/docs/local-rollout/reward-rubrics.md
index ca049114..24cdf0a1 100644
--- a/docs/local-rollout/reward-rubrics.md
+++ b/docs/local-rollout/reward-rubrics.md
@@ -15,7 +15,7 @@ def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -
return float_score
```
-> **Note:** The runtime forwards `None` for `ground_truth` when no reference answer exists. Annotate the parameter as `Optional[str]` (or handle `None` explicitly) if your rubric logic expects to run in that scenario.
+> **Note:** The runtime forwards `None` for `ground_truth` when no reference answer exists.
### Required `extra_info` Fields
@@ -30,22 +30,16 @@ def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -
- **`score_min` / `score_max`** -- Optional numeric overrides for the expected score range.
- **`model_info_overrides`** -- Optional dict merged into the provider configuration.
-Additional keys are passthrough and can be used for custom configuration. The decorator enforces the parameter names/annotations, validates the embedded configuration at call time, and ensures the wrapped function returns a `float`.
+> **Note:** `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
-> **Note:** Annotation quirk: `extra_info` must be annotated as `dict` **without** a default value, unlike `@osmosis_reward`.
+## evaluate_rubric()
-## evaluate_rubric
-
-`evaluate_rubric` talks to hosted LLM providers through [LiteLLM](https://github.com/BerriAI/litellm), a unified interface supporting 100+ providers. Every provider returns a strict JSON object with `{"score": number, "explanation": string}`.
-
-### Basic Usage
+Evaluate a solution against a natural-language rubric using an external LLM judge via [LiteLLM](https://github.com/BerriAI/litellm).
```python
from osmosis_ai import evaluate_rubric
-solution = "The capital of France is Paris."
-
-rubric_score = evaluate_rubric(
+score = evaluate_rubric(
rubric="Assistant must mention the verified capital city.",
solution_str=solution,
model_info={
@@ -54,35 +48,42 @@ rubric_score = evaluate_rubric(
"api_key_env": "OPENAI_API_KEY",
},
ground_truth="Paris",
+ metadata=None, # optional structured context quoted in the judge prompt
+ return_details=False, # True returns RewardRubricRunResult with score, explanation, raw
)
-
-print(rubric_score) # -> 1.0 (full payload available via return_details=True)
```
-### Credentials
+When `return_details=True`, returns a `RewardRubricRunResult` dict:
-Credentials are resolved from environment variables by default:
+```python
+class RewardRubricRunResult(TypedDict):
+ score: float # The clamped rubric score
+ explanation: str # The judge model's explanation
+ raw: Any # Full raw response from the LLM provider
+```
-- `OPENAI_API_KEY` for OpenAI
-- `ANTHROPIC_API_KEY` for Anthropic
-- `GEMINI_API_KEY` for Google Gemini
-- `XAI_API_KEY` for xAI
-- `OPENROUTER_API_KEY` for OpenRouter
-- `CEREBRAS_API_KEY` for Cerebras
+### Credentials
-Override the environment variable name with `model_info={"api_key_env": "CUSTOM_ENV_NAME"}` when needed, or supply an inline secret with `model_info={"api_key": "sk-..."}` for ephemeral credentials.
+Credentials are resolved in order: `api_key` (direct) > `api_key_env` (env var name) > provider default:
-`api_key` and `api_key_env` are mutually exclusive. When `api_key` is present and non-empty it is used directly, skipping any environment lookup. Otherwise the resolver falls back to `api_key_env` (or the provider default) and pulls the value from your local environment with `os.getenv`.
+| Provider | Default Environment Variable |
+|----------|------------------------------|
+| OpenAI | `OPENAI_API_KEY` |
+| Anthropic | `ANTHROPIC_API_KEY` |
+| Google Gemini | `GEMINI_API_KEY` |
+| xAI | `XAI_API_KEY` |
+| OpenRouter | `OPENROUTER_API_KEY` |
+| Cerebras | `CEREBRAS_API_KEY` |
-## Provider Examples
+Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
-### OpenAI
+## Example
```python
from osmosis_ai import osmosis_rubric, evaluate_rubric
@osmosis_rubric
-def compute_rubric_score_openai(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
+def compute_rubric_score(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
return evaluate_rubric(
rubric="Evaluate whether the solution correctly matches the expected answer.",
solution_str=solution_str,
@@ -95,40 +96,19 @@ def compute_rubric_score_openai(solution_str: str, ground_truth: str | None, ext
)
```
-### Anthropic
+## Error Types
-```python
-@osmosis_rubric
-def compute_rubric_score_anthropic(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
- return evaluate_rubric(
- rubric="Evaluate whether the solution correctly matches the expected answer.",
- solution_str=solution_str,
- ground_truth=ground_truth,
- model_info={
- "provider": "anthropic",
- "model": "claude-sonnet-4-5-20250929",
- "api_key_env": "ANTHROPIC_API_KEY",
- },
- )
-```
-
-Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
+| Exception | Description |
+|-----------|-------------|
+| `MissingAPIKeyError` | Required API key not found in environment. Message explains which env var to export. |
+| `ProviderRequestError` | LLM provider returned an error (auth, rate limit, timeout, invalid response). |
+| `ModelNotFoundError` | Model identifier not recognized by the provider. Subclass of `ProviderRequestError`. |
## File Placement
-In a Local Rollout repository, place rubric functions in the `reward_rubric/` directory:
-
-```
-reward_rubric/
-├── reward_rubric_openai.py
-├── reward_rubric_anthropic.py
-└── reward_rubric_xai.py
-```
-
-Osmosis discovers and syncs all `@osmosis_rubric`-decorated functions from this directory. See the [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) for working examples.
+In a Local Rollout repository, place rubric functions in the `reward_rubric/` directory. Osmosis discovers and syncs all `@osmosis_rubric`-decorated functions from this directory. See the [osmosis-git-sync-example](https://github.com/Osmosis-AI/osmosis-git-sync-example) for working examples.
## See Also
- [Reward Functions](./reward-functions.md) -- deterministic scoring with `@osmosis_reward`
-- [Rewards API Reference](../rewards-api.md) -- full API reference for decorators and `evaluate_rubric`
- [Local Rollout Overview](./overview.md) -- repository structure and setup
diff --git a/docs/remote-rollout/agent-loop.md b/docs/remote-rollout/agent-loop.md
index 23a7a39d..c66be890 100644
--- a/docs/remote-rollout/agent-loop.md
+++ b/docs/remote-rollout/agent-loop.md
@@ -27,33 +27,23 @@ class CalculatorAgentLoop(RolloutAgentLoop):
|-----------|------|-------------|
| `name` | `str` | Unique identifier for this agent loop (required) |
-Concrete subclasses must define `name`; omitting it raises `TypeError` at class definition time. Abstract subclasses (with remaining abstract methods) are not validated, so intermediate base classes don't need `name`.
+Concrete subclasses must define `name`; omitting it raises `TypeError` at class definition time.
**Abstract Methods:**
#### `get_tools(request: RolloutRequest) -> list[OpenAIFunctionToolSchema]`
-Return tools available for this rollout.
-
-- **Parameters:**
- - `request`: The incoming rollout request
-- **Returns:** List of `OpenAIFunctionToolSchema` objects
-- **Notes:** Can return different tools based on `request.metadata`
+Return tools available for this rollout. Can return different tools based on `request.metadata`.
#### `async run(ctx: RolloutContext) -> RolloutResult`
-Execute the agent loop.
-
-- **Parameters:**
- - `ctx`: Execution context with LLM client and request
-- **Returns:** `RolloutResult` with final status and messages
-- **Notes:** Messages must be append-only; never modify previous messages
+Execute the agent loop. Messages must be append-only; never modify previous messages.
**Convenience Methods:**
#### `get_default_tools() -> list[OpenAIFunctionToolSchema]`
-Return the default tool list for discovery and validation purposes, without requiring a real `RolloutRequest`. Calls `get_tools()` with a synthetic request. Used by the validation framework and platform registration.
+Return the default tool list for discovery and validation. Calls `get_tools()` with a synthetic request.
---
@@ -69,34 +59,16 @@ class RolloutContext:
llm: LLMClientProtocol
```
-**Attributes:**
-
-| Attribute | Type | Description |
-|-----------|------|-------------|
-| `request` | `RolloutRequest` | Original rollout request |
-| `tools` | `list[OpenAIFunctionToolSchema]` | Tools returned by `get_tools()` |
-| `llm` | `LLMClientProtocol` | LLM client (production: `OsmosisLLMClient` calling TrainGate; test/eval: `ExternalLLMClient` via LiteLLM) |
-
**Methods:**
#### `async chat(messages, **kwargs) -> CompletionsResult`
Shorthand for `self.llm.chat_completions()`.
-```python
-result = await ctx.chat(messages, temperature=0.7)
-```
-
#### `complete(final_messages, finish_reason="stop", reward=None) -> RolloutResult`
Create a successful completion result.
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `final_messages` | `List[Dict[str, Any]]` | required | Final conversation messages |
-| `finish_reason` | `str` | `"stop"` | Why the rollout ended |
-| `reward` | `float \| None` | `None` | Optional precomputed trajectory reward score |
-
#### `error(error_message, final_messages=None) -> RolloutResult`
Create an error result.
@@ -109,17 +81,6 @@ Record a tool call for metrics.
Log a debug event to the rollout's JSONL file. No-op if debug logging is not enabled.
-```python
-ctx.log_event("pre_llm", turn=0, num_messages=len(messages))
-ctx.log_event("llm_response", turn=0, has_tool_calls=result.has_tool_calls)
-```
-
-**Properties:**
-
-| Property | Type | Description |
-|----------|------|-------------|
-| `debug_enabled` | `bool` | Whether debug logging is enabled |
-
---
### RolloutResult
@@ -139,72 +100,6 @@ class RolloutResult:
---
-## Client
-
-### OsmosisLLMClient
-
-HTTP client for calling TrainGate's LLM completion endpoint.
-
-```python
-from osmosis_ai.rollout import OsmosisLLMClient
-
-async with OsmosisLLMClient(
- server_url="http://trainer:8080",
- rollout_id="rollout-123",
- timeout_seconds=300.0,
- max_retries=3,
-) as client:
- result = await client.chat_completions(messages)
-```
-
-**Constructor Parameters:**
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `server_url` | `str` | required | TrainGate base URL |
-| `rollout_id` | `str` | required | Unique rollout identifier |
-| `api_key` | `Optional[str]` | `None` | Bearer token for TrainGate callbacks (provided by TrainGate in `RolloutRequest.api_key`) |
-| `timeout_seconds` | `float` | `300.0` | Request timeout (seconds) |
-| `max_retries` | `int` | `3` | Max retries for transient errors (5xx / timeouts / transport) |
-| `complete_rollout_retries` | `int` | `2` | Max retries for `/v1/rollout/completed` callback |
-| `settings` | `Optional[RolloutClientSettings]` | `None` | Advanced client settings |
-
-**Methods:**
-
-#### `async chat_completions(messages, **kwargs) -> CompletionsResult`
-
-Call TrainGate's `/v1/chat/completions` endpoint.
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `messages` | `List[Dict]` | required | Conversation history |
-| `temperature` | `float` | `1.0` | Sampling temperature |
-| `top_p` | `float` | `1.0` | Top-p sampling |
-| `max_tokens` | `int` | `512` | Maximum response tokens |
-| `stop` | `Optional[Union[str, List[str]]]` | `None` | Stop sequences |
-| `logprobs` | `bool` | `True` | Return log probabilities |
-
-Extra/unknown keyword arguments are accepted and ignored.
-
-#### `async complete_rollout(status, final_messages, ...) -> None`
-
-Notify TrainGate that rollout is complete.
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `status` | `str` | required | "COMPLETED" or "ERROR" |
-| `final_messages` | `List[Dict]` | required | Final conversation |
-| `finish_reason` | `str` | `"stop"` | Why rollout ended |
-| `error_message` | `Optional[str]` | `None` | Error description |
-| `metrics` | `Optional[RolloutMetrics]` | `None` | Execution metrics |
-| `reward` | `Optional[float]` | `None` | Precomputed trajectory reward score |
-
-#### `get_metrics() -> RolloutMetrics`
-
-Get current metrics from this client session.
-
----
-
### CompletionsResult
Result from LLM completion call.
@@ -241,26 +136,6 @@ from osmosis_ai.rollout import serve_agent_loop
serve_agent_loop(agent_loop, port=9000)
```
-**Parameters:**
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `agent_loop` | `RolloutAgentLoop` | required | Your agent implementation |
-| `host` | `str` | `"0.0.0.0"` | Host to bind to |
-| `port` | `int` | `9000` | Port to bind to |
-| `validate` | `bool` | `True` | Validate agent loop before starting |
-| `log_level` | `str` | `"info"` | Uvicorn log level |
-| `reload` | `bool` | `False` | Enable auto-reload for development |
-| `settings` | `Optional[RolloutSettings]` | `None` | Configuration override |
-| `skip_register` | `bool` | `False` | Skip registering with Osmosis Platform |
-| `api_key` | `Optional[str]` | `None` | RolloutServer API key used by TrainGate (generated if not provided) |
-| `local_debug` | `bool` | `False` | Disable API key auth and force `skip_register=True` |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging |
-
-**Raises:** `ImportError` (missing FastAPI/uvicorn), `AgentLoopValidationError` (validation fails)
-
----
-
### create_app()
Factory function to create a FastAPI application.
@@ -271,29 +146,13 @@ from osmosis_ai.rollout import create_app
app = create_app(agent_loop, max_concurrent=100, record_ttl_seconds=3600)
```
-**Parameters:**
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `agent_loop` | `RolloutAgentLoop` | required | Your agent implementation |
-| `max_concurrent` | `int \| None` | `None` | Max concurrent rollouts |
-| `record_ttl_seconds` | `float \| None` | `None` | TTL for completed records |
-| `settings` | `Optional[RolloutSettings]` | `None` | Global settings override |
-| `credentials` | `Optional[WorkspaceCredentials]` | `None` | Workspace credentials for platform registration |
-| `server_host` | `Optional[str]` | `None` | Host (for platform registration) |
-| `server_port` | `Optional[int]` | `None` | Port (for platform registration) |
-| `api_key` | `Optional[str]` | `None` | API key for authenticating incoming requests |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logging |
-| `on_startup` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async startup callback |
-| `on_shutdown` | `Optional[Callable[[], Awaitable[None]]]` | `None` | Async shutdown callback |
-
-**Generated Endpoints:**
-
-| Endpoint | Method | Status | Description |
-|----------|--------|--------|-------------|
-| `/v1/rollout/init` | POST | 202 | Accept rollout request |
-| `/health` | GET | 200 | Health check (unauthenticated) |
-| `/platform/health` | GET | 200 | Authenticated health check for Osmosis Platform |
+See docstrings for full parameter lists. Generated endpoints:
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/v1/rollout/init` | POST | Accept rollout request (returns 202) |
+| `/health` | GET | Health check |
+| `/platform/health` | GET | Authenticated health check for Osmosis Platform |
---
@@ -301,8 +160,6 @@ app = create_app(agent_loop, max_concurrent=100, record_ttl_seconds=3600)
### validate_agent_loop()
-Validate a RolloutAgentLoop implementation.
-
```python
from osmosis_ai.rollout import validate_agent_loop
@@ -311,35 +168,7 @@ if not result.valid:
result.raise_if_invalid()
```
-**Parameters:**
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `agent_loop` | `RolloutAgentLoop` | required | Agent loop to validate |
-| `request` | `Optional[RolloutRequest]` | `None` | Custom request for testing get_tools() |
-
-**Returns:** `ValidationResult` with `valid`, `errors`, `warnings`, `agent_name`, `tool_count`.
-
-**Validation Checks:**
-- `name` attribute is defined and non-empty
-- `get_tools()` returns a valid list
-- Each tool conforms to OpenAI function schema
-- `run()` method is async
-
-**Common Error Codes:** `MISSING_NAME`, `EMPTY_NAME`, `GET_TOOLS_EXCEPTION`, `GET_TOOLS_RETURNS_NONE`, `MISSING_TOOL_TYPE`, `MISSING_FUNCTION`, `MISSING_FUNCTION_NAME`, `RUN_NOT_ASYNC`
-
----
-
-## Registry
-
-Functions for managing multiple agent loop instances:
-
-| Function | Description |
-|----------|-------------|
-| `register_agent_loop(instance)` | Register an agent loop. Raises `ValueError` if name is already registered. |
-| `get_agent_loop(name)` | Get by name. Raises `AgentLoopNotFoundError` if not found. |
-| `list_agent_loops()` | List all registered agent loop names. |
-| `unregister_agent_loop(name)` | Remove from registry. Returns `True` if found. |
+Returns `ValidationResult` with `valid`, `errors`, `warnings`, `agent_name`, `tool_count`.
---
@@ -347,8 +176,6 @@ Functions for managing multiple agent loop instances:
### RolloutRequest
-Request sent to `/v1/rollout/init` to start a rollout.
-
```python
class RolloutRequest(BaseModel):
rollout_id: str # 1-256 chars, not empty
@@ -358,15 +185,13 @@ class RolloutRequest(BaseModel):
tool_server_url: Optional[str] = None
max_turns: int = 10
max_tokens_total: int = 8192
- metadata: Dict[str, Any] = {} # JSON-serializable (size-limited; default 1MB)
- api_key: Optional[str] = None # Optional Bearer token for callbacks to TrainGate
+ metadata: Dict[str, Any] = {}
+ api_key: Optional[str] = None
idempotency_key: Optional[str] = None
```
### InitResponse
-Response from `/v1/rollout/init` endpoint.
-
```python
class InitResponse(BaseModel):
rollout_id: str
@@ -375,8 +200,6 @@ class InitResponse(BaseModel):
### RolloutResponse
-Posted to `/v1/rollout/completed`.
-
```python
class RolloutResponse(BaseModel):
rollout_id: str
@@ -384,15 +207,13 @@ class RolloutResponse(BaseModel):
final_messages: List[MessageDict] = []
finish_reason: Optional[str] = None
error_message: Optional[str] = None
- reward: Optional[float] = None # Precomputed trajectory reward
+ reward: Optional[float] = None
metrics: Optional[RolloutMetrics] = None
extra_fields: Dict[str, Any] = {}
```
### OpenAIFunctionToolSchema
-OpenAI-compatible tool definition.
-
```python
class OpenAIFunctionToolSchema(BaseModel):
type: str # "function"
@@ -401,24 +222,12 @@ class OpenAIFunctionToolSchema(BaseModel):
class OpenAIFunctionSchema(BaseModel):
name: str
description: str
- parameters: OpenAIFunctionParametersSchema # defaults to empty object
+ parameters: OpenAIFunctionParametersSchema
strict: bool = False
-
-class OpenAIFunctionParametersSchema(BaseModel):
- type: str # usually "object"
- properties: Dict[str, OpenAIFunctionPropertySchema]
- required: List[str]
-
-class OpenAIFunctionPropertySchema(BaseModel):
- type: str
- description: Optional[str] = None
- enum: Optional[List[str]] = None
```
### RolloutMetrics
-Execution metrics.
-
```python
class RolloutMetrics(BaseModel):
total_latency_ms: float = 0.0
@@ -431,14 +240,6 @@ class RolloutMetrics(BaseModel):
max_context_tokens: int = 0
```
-### RolloutStatus
-
-```python
-class RolloutStatus(str, Enum):
- COMPLETED = "COMPLETED"
- ERROR = "ERROR"
-```
-
---
## Exceptions
@@ -449,35 +250,15 @@ All exceptions inherit from `OsmosisRolloutError`.
|-----------|-------------|
| `OsmosisRolloutError` | Base exception for all rollout errors |
| `OsmosisTransportError` | Network/transport level errors |
-| `OsmosisServerError` | Server returned 5xx (retryable). Has `status_code` attribute. |
-| `OsmosisValidationError` | Server returned 4xx (not retryable). Has `status_code` attribute. |
+| `OsmosisServerError` | Server returned 5xx (retryable). Has `status_code`. |
+| `OsmosisValidationError` | Server returned 4xx (not retryable). Has `status_code`. |
| `OsmosisTimeoutError` | Request timed out |
-| `AgentLoopNotFoundError` | Agent loop not found in registry. Has `name` and `available` attributes. |
+| `AgentLoopNotFoundError` | Agent loop not found in registry |
| `AgentLoopValidationError` | Validation failed. Has `errors` attribute. |
-| `ToolExecutionError` | Tool execution failed. Has optional `tool_call_id` and `tool_name`. |
-| `ToolArgumentError` | Tool argument parsing failed (inherits from `ToolExecutionError`). |
-
----
-
-## Type Aliases
-
-### MessageDict
-
-```python
-MessageDict = Dict[str, Any]
-# Example: {"role": "user", "content": "Hello"}
-# Example: {"role": "tool", "content": "42", "tool_call_id": "call_123"}
-```
-
-### SamplingParamsDict
-
-```python
-SamplingParamsDict = Dict[str, Any]
-# Example: {"temperature": 0.7, "max_tokens": 512}
-```
+| `ToolExecutionError` | Tool execution failed |
+| `ToolArgumentError` | Tool argument parsing failed (inherits `ToolExecutionError`) |
## See Also
-- [Remote Rollout Overview](./overview.md) - Quick start guide
- [Architecture](./architecture.md) - Protocol design
- [Examples](./examples.md) - Working code examples
diff --git a/docs/remote-rollout/architecture.md b/docs/remote-rollout/architecture.md
index 3d1a516e..3865ae0b 100644
--- a/docs/remote-rollout/architecture.md
+++ b/docs/remote-rollout/architecture.md
@@ -1,11 +1,9 @@
# Architecture
-This document describes the architecture and communication protocol of the Osmosis Remote Rollout SDK.
+The Remote Rollout system separates **LLM inference** (on training cluster) from **agent logic** (on your RolloutServer).
## System Overview
-The Remote Rollout system consists of two main components:
-
```
┌─────────────┐ ┌─────────────────┐
│ TrainGate │ ◄──── HTTP ────────► │ RolloutServer │
@@ -41,8 +39,6 @@ TrainGate sends a `RolloutRequest` to `/v1/rollout/init`. RolloutServer:
### 2. Agent Loop Execution
-The agent runs asynchronously after returning the 202 response:
-
```
RolloutServer TrainGate
│ │
@@ -71,8 +67,6 @@ Key points:
### 3. Completion Notification
-When the agent loop finishes:
-
```
RolloutServer TrainGate
│ │
@@ -86,148 +80,7 @@ RolloutServer TrainGate
│ │
```
-The `RolloutResponse` includes:
-- `status`: "COMPLETED" or "ERROR"
-- `final_messages`: Complete conversation history
-- `finish_reason`: Why the rollout ended
-- `metrics`: Execution statistics
-
-## Core Components
-
-### RolloutAgentLoop
-
-The abstract base class that users implement:
-
-```python
-class RolloutAgentLoop(ABC):
- name: str # Unique identifier
-
- @abstractmethod
- def get_tools(self, request: RolloutRequest) -> list[OpenAIFunctionToolSchema]:
- """Return tools for this rollout."""
-
- @abstractmethod
- async def run(self, ctx: RolloutContext) -> RolloutResult:
- """Execute the agent loop."""
-```
-
-The `__init_subclass__` hook validates that subclasses define `name`.
-
-### RolloutContext
-
-Execution context provided to the agent:
-
-```python
-@dataclass
-class RolloutContext:
- request: RolloutRequest # Original request
- tools: list[OpenAIFunctionToolSchema] # Available tools
- llm: LLMClientProtocol # LLM client (OsmosisLLMClient in production)
-
- async def chat(self, messages, **kwargs) -> CompletionsResult
- def complete(self, messages, finish_reason="stop", reward=None) -> RolloutResult
- def error(self, message, final_messages=None) -> RolloutResult
- def record_tool_call(self, latency_ms=0.0) -> None
-```
-
-### OsmosisLLMClient
-
-HTTP client for TrainGate communication:
-
-- Connection pooling (up to 100 connections)
-- Automatic retry with exponential backoff (5xx errors, timeouts, transport errors)
-- Timeout handling (default 300 seconds)
-- Metrics collection (latency, token counts)
-
-### AppState
-
-Server state management:
-
-- Tracks running and completed rollouts
-- Provides idempotency (duplicate requests return same response)
-- Background cleanup of old records
-- Concurrency limiting via semaphore
-
-## Error Handling
-
-### Exception Hierarchy
-
-```
-OsmosisRolloutError (base)
-├── OsmosisTransportError # Network errors
-├── OsmosisServerError # 5xx errors (retryable)
-├── OsmosisValidationError # 4xx errors (not retryable)
-├── OsmosisTimeoutError # Timeouts
-├── AgentLoopNotFoundError # Registry lookup failed
-├── ToolExecutionError # Tool execution errors
-└── ToolArgumentError # Tool argument parsing errors
-```
-
-### Retry Strategy
-
-| Error Type | Retried? | Notes |
-|------------|----------|-------|
-| 5xx Server Error | Yes | Exponential backoff (1s, 2s, 4s, ...) |
-| 4xx Client Error | No | Fails immediately |
-| Timeout | Yes | Up to max_retries |
-| Network Error | Yes | Connection failures, DNS, etc. |
-
-### Agent Error Handling
-
-If an exception occurs in `agent_loop.run()`:
-
-1. The error is logged
-2. `complete_rollout()` is called with status="ERROR"
-3. The error message is included in the response
-
-## Metrics Collection
-
-The SDK automatically tracks:
-
-| Metric | Description |
-|--------|-------------|
-| `total_latency_ms` | Total rollout duration |
-| `llm_latency_ms` | Time spent in LLM calls |
-| `tool_latency_ms` | Time spent in tool execution |
-| `num_llm_calls` | Number of LLM completion calls |
-| `num_tool_calls` | Number of tool executions |
-| `prompt_tokens` | Total prompt tokens |
-| `response_tokens` | Total response tokens |
-
-Use `ctx.record_tool_call(latency_ms)` to track tool execution time.
-
-## Idempotency
-
-The server handles duplicate requests:
-
-1. Check if `rollout_id` is already running or recently completed
-2. If duplicate, return the same `InitResponse` without starting new task
-3. Completed rollout records are kept for `record_ttl_seconds` (default: 1 hour)
-
-This ensures safety during network retries.
-
-## Concurrency Control
-
-The server limits concurrent rollouts via `max_concurrent` parameter:
-
-```python
-app = create_app(agent_loop, max_concurrent=100)
-```
-
-When the limit is reached, new rollouts wait for a slot to become available.
-
-## Lifecycle Management
-
-On server startup:
-- Background cleanup task starts (prunes old completed records every 60s)
-
-On server shutdown:
-- Cleanup task is cancelled
-- All running rollout tasks are cancelled
-- Resources are released
-
## See Also
-- [Remote Rollout Overview](./overview.md) -- Quick start guide
- [Agent Loop Guide](./agent-loop.md) -- Complete API documentation
- [Examples](./examples.md) -- Working code examples
diff --git a/docs/remote-rollout/examples.md b/docs/remote-rollout/examples.md
index 6c69f4e7..df8b1e43 100644
--- a/docs/remote-rollout/examples.md
+++ b/docs/remote-rollout/examples.md
@@ -1,10 +1,6 @@
# Examples
-Complete working examples for the Osmosis Remote Rollout SDK.
-
-
-Full source code is available in the [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) repository. The snippets below are abbreviated for clarity.
-
+Full source code is available in the [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) repository.
## CLI Quick Start
@@ -28,21 +24,6 @@ osmosis eval -m server:agent_loop -d test_data.jsonl \
The module path format is `module:attribute`. The CLI automatically adds the current directory to Python path.
----
-
-## Project Structure
-
-```
-rollout-server/
-├── server.py # Agent loop + FastAPI app
-├── tools.py # Tool definitions and execution
-├── rewards.py # Reward computation
-├── test_data.jsonl # Test dataset
-└── pyproject.toml # Dependencies: osmosis-ai[server]>=0.2.14
-```
-
----
-
## Debug Logging
Enable debug logging with the `--log` CLI flag or `ROLLOUT_DEBUG_DIR` environment variable. Use `ctx.log_event()` in your agent's `run()` method to record events -- these are no-ops unless debug logging is enabled.
@@ -57,41 +38,6 @@ ROLLOUT_DEBUG_DIR=./rollout_logs osmosis serve -m server:agent_loop
Each server session creates a timestamped subdirectory, and each rollout creates a JSONL file with events like `pre_llm`, `llm_response`, `tool_results`, and `rollout_complete`.
----
-
-## Tool Utilities
-
-The `osmosis_ai.rollout.tools` module provides utilities for tool call execution:
-
-| Function | Description |
-|----------|-------------|
-| `get_tool_call_info(tool_call)` | Extract `(call_id, name, args)` from a tool call dict, parsing JSON arguments |
-| `serialize_tool_result(value)` | Convert any value to a string suitable for tool results |
-| `create_tool_result(call_id, content)` | Create a `{"role": "tool", ...}` message dict |
-| `create_tool_error_result(call_id, error)` | Create an error tool result message |
-| `execute_tool_calls(calls, executor)` | Execute multiple tool calls concurrently via an async executor function |
-
-See the example repo's `tools.py` for the recommended pattern using these utilities.
-
----
-
-## Message Utilities
-
-The `osmosis_ai.rollout` module exports helpers for working with messages:
-
-| Function | Description |
-|----------|-------------|
-| `parse_tool_calls(message)` | Safely extract `tool_calls` from an assistant message |
-| `normalize_stop(value)` | Normalize stop parameter to `List[str]` or `None` |
-| `get_message_content(message)` | Get message content safely |
-| `is_assistant_message(message)` | Check if message role is `assistant` |
-| `is_tool_message(message)` | Check if message role is `tool` |
-| `is_user_message(message)` | Check if message role is `user` |
-| `count_messages_by_role(messages)` | Count messages by role, returns `{"user": N, ...}` |
-
----
-
## See Also
-- [Testing](./testing.md) -- unit tests and mock trainer
- [Agent Loop Guide](./agent-loop.md) -- endpoints, schemas, types
diff --git a/docs/remote-rollout/overview.md b/docs/remote-rollout/overview.md
index fee5318f..eaf82eeb 100644
--- a/docs/remote-rollout/overview.md
+++ b/docs/remote-rollout/overview.md
@@ -2,49 +2,10 @@
A lightweight SDK for integrating agent frameworks with the Osmosis remote rollout protocol. This SDK handles protocol communication between your agent logic and the Osmosis training infrastructure.
-## Overview
-
-The Remote Rollout SDK enables you to:
-
-- Define custom agent loops that integrate with Osmosis training
-- Handle LLM completions through the TrainGate service
-- Execute tools and collect training data (trajectories, token IDs, logprobs)
-- Run agents as scalable HTTP servers
-
For an alternative approach that requires no server infrastructure, see [Local Rollout](../local-rollout/overview.md).
-## Installation
-
-```bash
-# Basic installation
-pip install osmosis-ai
-
-# With server support (FastAPI + uvicorn)
-pip install osmosis-ai[server]
-
-# Full installation with all optional features
-pip install osmosis-ai[full]
-```
-
## Quick Start
-### Project Structure
-
-```
-rollout-server/
-├── server.py # Agent loop + FastAPI app
-├── tools.py # Tool definitions and execution
-├── rewards.py # Reward computation
-├── test_data.jsonl # Test dataset
-└── pyproject.toml # Dependencies: osmosis-ai[server]>=0.2.14
-```
-
-For the complete working project, see [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example).
-
-### Step 1: Implement Your Agent Loop
-
-Create `server.py` with a class that inherits from `RolloutAgentLoop`:
-
```python
class CalculatorAgentLoop(RolloutAgentLoop):
name = "calculator"
@@ -67,89 +28,21 @@ agent_loop = CalculatorAgentLoop()
app = create_app(agent_loop)
```
-See [Examples](./examples.md) for the complete `tools.py` and `rewards.py` files.
-
-### Step 2: Run the Server
-
```bash
-# Validate agent loop (checks tools, async run method, etc.)
+# Validate agent loop
osmosis validate -m server:agent_loop
-# Start server with validation (default port 9000)
+# Start server (default port 9000)
osmosis serve -m server:agent_loop
-# Specify port and enable auto-reload
-osmosis serve -m server:agent_loop -p 8080 --reload
-```
-
-For programmatic alternatives (`create_app()`, `serve_agent_loop()`), see [Agent Loop Guide](./agent-loop.md#server).
-
-## Key Concepts
-
-### RolloutAgentLoop
-
-The abstract base class for your agent implementation. You must:
-
-- Set the `name` class attribute (unique identifier)
-- Implement `get_tools()` to return available tools
-- Implement `run()` to execute your agent logic
-
-### RolloutContext
-
-Provided to your `run()` method, containing:
-
-- `request`: The original `RolloutRequest` with messages and params
-- `tools`: List of tools returned by `get_tools()`
-- `llm`: The `OsmosisLLMClient` for LLM calls
-
-Key methods: `ctx.chat()`, `ctx.complete()`, `ctx.error()`, `ctx.record_tool_call()`
-
-### RolloutResult
-
-The return value from your agent loop:
-
-- `status`: "COMPLETED" or "ERROR"
-- `final_messages`: The complete conversation history
-- `finish_reason`: Why the rollout ended
-- `metrics`: Execution metrics (latency, token counts)
-
-## Server Endpoints
-
-The server created by `create_app()` exposes:
-
-| Endpoint | Method | Description |
-|----------|--------|-------------|
-| `/v1/rollout/init` | POST | Accept rollout requests (returns 202) |
-| `/health` | GET | Health check |
-
-## Configuration Options
-
-```python
-app = create_app(
- agent_loop,
- max_concurrent=100, # Max concurrent rollouts
- record_ttl_seconds=3600, # How long to keep completed records
-)
+# Test locally with a cloud LLM
+osmosis test -m server:agent_loop -d test_data.jsonl --model gpt-5-mini
```
-## Settings (Environment Variables)
-
-When installed with `osmosis-ai[server]`, settings can be loaded from environment variables (or a `.env` file) using these prefixes:
-
-- `OSMOSIS_ROLLOUT_CLIENT_*` - Client settings (timeout, retries, etc.)
-- `OSMOSIS_ROLLOUT_SERVER_*` - Server settings (concurrency, TTL, etc.)
-
-See [Configuration](../configuration.md) for the full list of environment variables.
-
-## Example Repository
-
-See the complete working example: [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example)
+For the complete working project, see [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example).
## Next Steps
-- [Architecture](./architecture.md) -- Understand the system design
+- [Architecture](./architecture.md) -- Protocol design and flow diagrams
- [Agent Loop Guide](./agent-loop.md) -- Complete API documentation
-- [Examples](./examples.md) -- Full working examples
-- [Testing](./testing.md) -- Unit tests and mock trainer
-- [Dataset Format](../datasets.md) -- Supported formats and required columns
-- [Test Mode](../test-mode.md) -- Local testing with external LLM providers
+- [Examples](./examples.md) -- CLI usage and debug logging
diff --git a/docs/remote-rollout/testing.md b/docs/remote-rollout/testing.md
deleted file mode 100644
index 5aa13d28..00000000
--- a/docs/remote-rollout/testing.md
+++ /dev/null
@@ -1,105 +0,0 @@
-# Testing Your Agent
-
-This guide covers unit testing and local testing strategies for rollout agents. These techniques let you validate your agent logic without deploying to training infrastructure or connecting to external LLM providers.
-
-## Using Mock Trainer
-
-The SDK provides a mock trainer for local testing without a real TrainGate server.
-
-```python
-# test_with_mock_trainer.py
-
-import pytest
-from fastapi.testclient import TestClient
-
-from osmosis_ai.rollout.testing import (
- create_mock_trainer_app,
- RolloutCompletionTracker,
- patch_httpx_for_mock_trainer,
-)
-
-
-@pytest.fixture
-def mock_trainer(monkeypatch):
- """Set up mock trainer with completion tracking."""
- tracker = RolloutCompletionTracker()
- app = create_mock_trainer_app(tracker=tracker)
- client = TestClient(app)
- patch_httpx_for_mock_trainer(client, monkeypatch)
- return client, tracker
-
-
-def test_rollout_with_mock_trainer(mock_trainer):
- """Test complete rollout flow with mock trainer."""
- client, tracker = mock_trainer
-
- # Your rollout will use the mock trainer
- # ...
-
- # Wait for completion callback
- assert tracker.wait(timeout=5.0)
- assert len(tracker.responses) == 1
- assert tracker.responses[0]["status"] == "COMPLETED"
-```
-
-## Custom Tool Call Generator
-
-You can customize when the mock trainer generates tool calls:
-
-```python
-def weather_tool_generator(message):
- """Generate weather tool calls for weather-related messages."""
- if "weather" in message.get("content", "").lower():
- return [
- {
- "id": "call_weather",
- "type": "function",
- "function": {"name": "get_weather", "arguments": "{}"},
- }
- ]
- return None
-
-app = create_mock_trainer_app(tool_call_generator=weather_tool_generator)
-```
-
-## Testing Utilities Reference
-
-The `osmosis_ai.rollout.testing` module provides the following public utilities for testing agent loops without a real TrainGate server.
-
-### `create_mock_trainer_app(tracker=None, tool_call_generator=None)`
-
-Creates a mock trainer FastAPI application that implements the same HTTP endpoints as a real TrainGate server:
-
-- `POST /v1/chat/completions` -- returns deterministic LLM responses with fake token IDs
-- `POST /v1/rollout/completed` -- accepts completion callbacks
-- `GET /v1/rollout/completed/{rollout_id}` -- queries completed rollouts
-- `GET /health` -- health check
-
-By default, the mock generates tool calls when it detects calculator-related keywords (e.g. "add", "calculate", "multiply") in the user message. Pass a custom `tool_call_generator` function to override this behavior.
-
-### `RolloutCompletionTracker`
-
-Thread-safe tracker that captures `/v1/rollout/completed` callbacks during tests.
-
-| Attribute / Method | Description |
-|--------------------|-------------|
-| `event` | `threading.Event` that is set when a completion is received |
-| `responses` | List of captured completion response dicts |
-| `record(response)` | Record a response and signal the event |
-| `clear()` | Clear recorded responses and reset the event |
-| `wait(timeout=5.0)` | Block until a completion is received; returns `True` on success, `False` on timeout |
-
-### `patch_httpx_for_mock_trainer(client, monkeypatch)`
-
-Patches `httpx.AsyncClient.post` so that any requests to `/v1/chat/completions` or `/v1/rollout/completed` are routed to the mock trainer `TestClient` instead of making real HTTP calls. All other requests pass through unchanged.
-
-### `fake_token_ids(text)` / `fake_prompt_token_ids(messages)`
-
-Generate deterministic fake token IDs for testing. These produce deterministic output suitable for snapshot testing but do not correspond to any real tokenizer.
-
-For a complete test file example, see the [osmosis-remote-rollout-example](https://github.com/Osmosis-AI/osmosis-remote-rollout-example) repository.
-
-## See Also
-
-- [Examples](./examples.md) -- agent implementations and utilities
-- [Test Mode](../test-mode.md) -- testing with cloud LLMs
diff --git a/docs/rewards-api.md b/docs/rewards-api.md
deleted file mode 100644
index cba1a507..00000000
--- a/docs/rewards-api.md
+++ /dev/null
@@ -1,222 +0,0 @@
-# Rewards API Reference
-
-API reference for reward function and rubric evaluation decorators. These are used in both Local Rollout and Remote Rollout modes for scoring LLM outputs during training.
-
-## @osmosis_reward
-
-Decorator for deterministic reward functions. Validates the function signature at decoration time and ensures a `float` return value at call time.
-
-### Required Signature
-
-```python
-from osmosis_ai import osmosis_reward
-
-@osmosis_reward
-def your_function(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
- return float_score
-```
-
-### Parameters
-
-| Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
-| `solution_str` | `str` | Yes | The solution string to evaluate |
-| `ground_truth` | `str` | Yes | The correct/expected answer |
-| `extra_info` | `dict` | No (default `None`) | Optional dictionary for additional configuration |
-| `**kwargs` | - | Yes | Required for platform compatibility |
-
-### Return Value
-
-Must return a `float`. The decorator raises `TypeError` if the return type check fails.
-
-### Validation Rules
-
-- Function must accept exactly `solution_str`, `ground_truth`, `extra_info` as positional parameters
-- `extra_info` must have a default value of `None`
-- Must include `**kwargs` for platform compatibility
-- Return value must be a `float` (checked at call time)
-
----
-
-## @osmosis_rubric
-
-Decorator for LLM-as-judge rubric functions. Validates the function signature and ensures required `extra_info` fields are present at call time.
-
-### Required Signature
-
-```python
-from osmosis_ai import osmosis_rubric
-
-@osmosis_rubric
-def your_rubric(solution_str: str, ground_truth: str | None, extra_info: dict) -> float:
- return float_score
-```
-
-### Parameters
-
-| Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
-| `solution_str` | `str` | Yes | The solution string to evaluate |
-| `ground_truth` | `str \| None` | Yes | The correct/expected answer (may be `None`) |
-| `extra_info` | `dict` | Yes (no default) | Configuration dict with provider, model, rubric |
-
-### Required `extra_info` Fields
-
-| Field | Type | Description |
-|-------|------|-------------|
-| `provider` | `str` | Non-empty string identifying the judge provider |
-| `model` | `str` | Non-empty string naming the provider model |
-| `rubric` | `str` | Natural-language rubric instructions |
-| `api_key` or `api_key_env` | `str` | Raw API key or environment variable name |
-
-### Optional `extra_info` Fields
-
-| Field | Type | Description |
-|-------|------|-------------|
-| `system_prompt` | `str` | Prepended to the provider's base system prompt |
-| `score_min` | `float` | Override minimum score bound |
-| `score_max` | `float` | Override maximum score bound |
-| `model_info_overrides` | `dict` | Merged into provider configuration |
-
-### Validation Rules
-
-- `extra_info` must be annotated as `dict` **without** a default value
-- `ground_truth` may be `None` (annotate as `Optional[str]` if needed)
-- Return type must be `float`
-
----
-
-## evaluate_rubric()
-
-Evaluate a solution against a natural-language rubric using an external LLM judge via [LiteLLM](https://github.com/BerriAI/litellm).
-
-### Signature
-
-```python
-from osmosis_ai import evaluate_rubric
-
-score = evaluate_rubric(
- rubric: str,
- solution_str: str,
- model_info: dict,
- ground_truth: str = None,
- metadata: dict = None,
- return_details: bool = False,
-)
-```
-
-### Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `rubric` | `str` | required | Natural-language rubric instructions for the judge |
-| `solution_str` | `str` | required | The solution to evaluate |
-| `model_info` | `dict` | required | Provider configuration (see below) |
-| `ground_truth` | `str` | `None` | Optional reference answer |
-| `metadata` | `dict` | `None` | Optional structured context quoted in the judge prompt |
-| `return_details` | `bool` | `False` | Return full `RewardRubricRunResult` payload |
-
-### `model_info` Fields
-
-| Field | Type | Required | Description |
-|-------|------|----------|-------------|
-| `provider` | `str` | Yes | Provider name (e.g., `"openai"`, `"anthropic"`) |
-| `model` | `str` | Yes | Model name (e.g., `"gpt-5"`, `"claude-sonnet-4-5-20250929"`) |
-| `api_key` | `str` | No | Raw API key (mutually exclusive with `api_key_env`) |
-| `api_key_env` | `str` | No | Environment variable name for API key |
-| `score_min` | `float` | No | Minimum score bound (default `0.0`) |
-| `score_max` | `float` | No | Maximum score bound (default `1.0`) |
-| `system_prompt` | `str` | No | Optional context prepended to judge prompt |
-| `original_input` | `str` | No | Optional original user input for context |
-| `timeout` | `float` | No | Provider timeout in seconds |
-| `reasoning_effort` | `str \| None` | No | Reasoning effort hint passed to the provider (e.g., `"low"`, `"medium"`, `"high"`). Silently dropped for models that do not support it. |
-
-### Return Value
-
-- When `return_details=False`: Returns a `float` score clamped to `[score_min, score_max]`
-- When `return_details=True`: Returns a `RewardRubricRunResult` dict with the following structure:
-
-```python
-class RewardRubricRunResult(TypedDict):
- score: float # The clamped rubric score
- explanation: str # The judge model's explanation for the score
- raw: Any # The full raw response from the LLM provider
-```
-
-### Credential Resolution
-
-Credentials are resolved in this order:
-
-1. `api_key` in `model_info` (used directly if present and non-empty)
-2. `api_key_env` in `model_info` (environment variable name to look up)
-3. Provider default environment variable:
-
-| Provider | Default Environment Variable |
-|----------|------------------------------|
-| OpenAI | `OPENAI_API_KEY` |
-| Anthropic | `ANTHROPIC_API_KEY` |
-| Google Gemini | `GEMINI_API_KEY` |
-| xAI | `XAI_API_KEY` |
-| OpenRouter | `OPENROUTER_API_KEY` |
-| Cerebras | `CEREBRAS_API_KEY` |
-
-### Provider Architecture
-
-All provider routing is handled by [LiteLLM](https://github.com/BerriAI/litellm). Any provider supported by LiteLLM can be used without additional configuration beyond setting the appropriate API key environment variable.
-
-Every provider returns a strict JSON object with `{"score": number, "explanation": string}`. The helper clamps the score into the configured range and validates the structure.
-
----
-
-## Error Types
-
-### MissingAPIKeyError
-
-Raised when the required API key is not found in the environment.
-
-```python
-from osmosis_ai import MissingAPIKeyError
-
-try:
- score = evaluate_rubric(...)
-except MissingAPIKeyError as e:
- print(f"Missing key: {e}")
- # Message explains which env var to export
-```
-
-### ProviderRequestError
-
-Raised when the LLM provider returns an error.
-
-```python
-from osmosis_ai import ProviderRequestError
-
-try:
- score = evaluate_rubric(...)
-except ProviderRequestError as e:
- print(f"Provider error: {e}")
-```
-
-### ModelNotFoundError
-
-Raised when the specified model identifier is not recognized by the provider. Subclass of `ProviderRequestError`.
-
-```python
-from osmosis_ai import ModelNotFoundError
-
-try:
- score = evaluate_rubric(...)
-except ModelNotFoundError as e:
- print(f"Model not found: {e}")
- # Check provider dashboard for latest model names
-```
-
-> Provider model snapshot names change frequently. Check each vendor's dashboard for the latest identifier if you encounter a "model not found" error.
-
----
-
-## See Also
-
-- [Reward Functions Guide](./local-rollout/reward-functions.md) -- usage guide with examples
-- [Reward Rubrics Guide](./local-rollout/reward-rubrics.md) -- usage guide with provider examples
-- [CLI Reference](./cli.md) -- `osmosis preview` and `osmosis eval-rubric` commands
diff --git a/docs/test-mode.md b/docs/test-mode.md
index 2af9f806..81430eb9 100644
--- a/docs/test-mode.md
+++ b/docs/test-mode.md
@@ -2,19 +2,6 @@
Test your agent implementations locally without TrainGate using external LLM providers via [LiteLLM](https://docs.litellm.ai/docs/providers). Works with both **Local Rollout** (MCP tools) and **Remote Rollout** (RolloutAgentLoop) agents.
-> **Note:** Breaking change: the legacy import path `osmosis_ai.rollout.test_mode` was removed.
-> Use `osmosis_ai.rollout.eval.common` and `osmosis_ai.rollout.eval.test_mode` instead.
-
-## Overview
-
-Test mode enables you to:
-
-- Validate agent logic before deploying to training infrastructure
-- **Two agent modes**: provide a `RolloutAgentLoop` with `-m`, or load MCP tools directly with `--mcp`
-- Use 100+ LLM providers (OpenAI, Anthropic, Groq, Ollama, etc.)
-- Run batch tests against datasets with detailed metrics
-- Debug step-by-step with interactive mode
-
## Quick Start
```bash
@@ -49,12 +36,6 @@ print(f"Passed: {results.passed}/{results.total}")
---
-## Dataset Format
-
-See [Dataset Format](./datasets.md) for supported formats and required columns.
-
----
-
## CLI Reference
```
@@ -120,7 +101,7 @@ See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for supported pr
## Interactive Mode
-Interactive mode allows step-by-step debugging of agent execution. After each LLM call, execution pauses to let you inspect the state.
+Interactive mode allows step-by-step debugging of agent execution.
### Commands
@@ -136,80 +117,6 @@ Interactive mode allows step-by-step debugging of agent execution. After each LL
---
-## API Reference
-
-### DatasetReader
-
-```python
-from osmosis_ai.rollout.eval.common import DatasetReader
-
-reader = DatasetReader("./data.jsonl")
-rows = reader.read(limit=10, offset=20)
-```
-
-| Method | Returns | Description |
-|--------|---------|-------------|
-| `read(limit, offset)` | `List[DatasetRow]` | Read rows with optional pagination |
-| `iter_rows()` | `Iterator[DatasetRow]` | Memory-efficient iterator |
-| `__len__()` | `int` | Total row count |
-
-`DatasetRow` is a `TypedDict` with `ground_truth`, `user_prompt`, `system_prompt`, plus any extra columns from the dataset.
-
-### ExternalLLMClient
-
-```python
-from osmosis_ai.rollout.eval.common import ExternalLLMClient
-
-client = ExternalLLMClient("gpt-5-mini") # OpenAI
-client = ExternalLLMClient("anthropic/claude-sonnet-4-5") # Anthropic
-client = ExternalLLMClient("ollama/llama3.1", api_base="http://localhost:11434") # Local
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `model` | `str` | required | Model name (simple or LiteLLM format) |
-| `api_key` | `Optional[str]` | `None` | API key (or use env var) |
-| `api_base` | `Optional[str]` | `None` | Base URL for OpenAI-compatible APIs |
-
-Method: `async chat_completions(messages, tools, **kwargs) -> CompletionsResult`
-
-### LocalTestRunner
-
-```python
-from osmosis_ai.rollout.eval.test_mode import LocalTestRunner
-
-runner = LocalTestRunner(agent_loop=agent, llm_client=client)
-result = await runner.run_single(row, row_index=0)
-batch_result = await runner.run_batch(rows=rows, max_turns=10)
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `agent_loop` | `RolloutAgentLoop` | required | Agent to test |
-| `llm_client` | `ExternalLLMClient` | required | LLM client |
-| `debug` | `bool` | `False` | Enable debug output |
-| `debug_dir` | `Optional[str]` | `None` | Directory for debug logs |
-
-### Result Types
-
-| Type | Key Fields |
-|------|------------|
-| `LocalTestRunResult` | `row_index`, `success`, `result: Optional[RolloutResult]`, `error`, `duration_ms`, `token_usage` |
-| `LocalTestBatchResult` | `results`, `total`, `passed`, `failed`, `total_duration_ms`, `total_tokens` |
-
-`token_usage` dict contains: `prompt_tokens`, `completion_tokens`, `total_tokens`, `num_llm_calls`.
-
-### InteractiveRunner
-
-```python
-from osmosis_ai.rollout.eval.test_mode import InteractiveRunner
-
-runner = InteractiveRunner(agent_loop=agent, llm_client=client)
-await runner.run_interactive_session(rows=rows, max_turns=10, initial_row=5)
-```
-
----
-
## Exceptions
All local workflow exceptions are shared by `test` and `eval` commands.
@@ -223,34 +130,8 @@ All local workflow exceptions are shared by `test` and `eval` commands.
---
-## Output Format
-
-When using `--output`, results are saved as JSON with this structure:
-
-- `summary` -- total, passed, failed, total_duration_ms, total_tokens
-- `results[]` -- per-row: row_index, success, error, duration_ms, token_usage, reward, finish_reason
-
----
-
-## Environment Variables
-
-API keys can be set via environment variables:
-
-| Provider | Environment Variable |
-|----------|---------------------|
-| OpenAI | `OPENAI_API_KEY` |
-| Anthropic | `ANTHROPIC_API_KEY` |
-| Groq | `GROQ_API_KEY` |
-| Azure | `AZURE_API_KEY` |
-
-See [LiteLLM Environment Variables](https://docs.litellm.ai/docs/providers) for more providers.
-
----
-
## See Also
- [Eval Mode](./eval-mode.md) -- Evaluate agents with eval functions and pass@k
- [Dataset Format](./datasets.md) -- Supported formats and required columns
-- [Remote Rollout Examples](./remote-rollout/examples.md) -- Working code examples
-- [Local Rollout MCP Tools](./local-rollout/mcp-tools.md) -- MCP tool definition
- [LiteLLM Providers](https://docs.litellm.ai/docs/providers) -- Supported LLM providers
diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md
index 1ebd841d..df8b26f9 100644
--- a/docs/troubleshooting.md
+++ b/docs/troubleshooting.md
@@ -119,24 +119,11 @@ numeric type.
## Rubric Evaluation Errors
-These errors originate from `osmosis_ai.rubric_eval` and
-`osmosis_ai.rubric_types`.
-
### MissingAPIKeyError
Raised when the LLM provider API key cannot be found. The error message
includes a hint showing the expected environment variable.
-Example:
-
-```
-MissingAPIKeyError: Environment variable 'OPENAI_API_KEY' is not set.
-Export it with your openai API key before calling evaluate_rubric.
-Set the required API key before running:
-
- export OPENAI_API_KEY="..."
-```
-
Resolution -- set the appropriate environment variable for your provider:
| Provider | Environment variable |
@@ -156,37 +143,17 @@ You can also provide the key directly in `model_info` via the `api_key` or
### ProviderRequestError
-Raised when the LLM provider call fails. The error includes the provider name,
-model name, and a detail string.
-
-```
-ProviderRequestError: Provider 'openai' request for model 'gpt-5-mini' failed.
-```
-
-Common causes:
+Raised when the LLM provider call fails. Common causes:
- **Authentication failure** -- your API key is invalid or expired.
-- **Rate limiting** -- you have exceeded the provider's rate limit. Wait and
- retry.
-- **Timeout** -- the request took too long. Try increasing the `timeout`
- parameter or use a faster model.
-- **Connection error** -- network issue between you and the provider. Check
- connectivity.
-- **Invalid JSON response** -- the model returned content that could not be
- parsed. Refine rubric instructions so the model returns valid JSON.
+- **Rate limiting** -- you have exceeded the provider's rate limit. Wait and retry.
+- **Timeout** -- the request took too long. Try increasing the `timeout` parameter or use a faster model.
+- **Invalid JSON response** -- the model returned content that could not be parsed.
### ModelNotFoundError
A subclass of `ProviderRequestError` raised when the requested model does not
-exist.
-
-```
-ModelNotFoundError: Provider 'openai' request for model 'gpt-99' failed.
-Model 'gpt-99' was not found. Confirm the model identifier is correct and your openai account has access to it.
-```
-
-Resolution -- verify the model name and that your account has access. Use the
-`provider/model` format, e.g. `openai/gpt-5-mini`, `anthropic/claude-sonnet-4-5`.
+exist. Verify the model name and that your account has access.
## Test Mode Errors
@@ -196,18 +163,6 @@ These errors occur when running `osmosis test` or `osmosis eval`.
Raised when the dataset file cannot be read or parsed.
-```
-DatasetParseError: Unsupported file format: .txt. Supported formats: .parquet (recommended), .jsonl, .csv
-```
-
-```
-DatasetParseError: Invalid JSON at line 5: Expecting ',' delimiter
-```
-
-```
-DatasetParseError: Parquet support requires pyarrow. Install with: pip install pyarrow
-```
-
Resolution:
- Ensure your file uses a supported format: `.parquet`, `.jsonl`, or `.csv`.
@@ -218,18 +173,6 @@ Resolution:
Raised when dataset rows are missing required columns or have invalid values.
-```
-DatasetValidationError: Row 0: Missing required columns: ['user_prompt']
-```
-
-```
-DatasetValidationError: Row 3: 'ground_truth' cannot be null
-```
-
-```
-DatasetValidationError: Row 7: 'system_prompt' must be a string, got int
-```
-
Every dataset row must include these columns:
| Column | Type | Description |
@@ -242,37 +185,21 @@ All values must be non-empty strings.
### ToolValidationError
-Raised when tool schemas returned by your agent are invalid for the provider's
-API.
-
-```
-ToolValidationError:
-```
-
-Resolution -- ensure your `get_tools()` method returns valid OpenAI-compatible
-function tool schemas. Each tool must have a `type` field set to `"function"` and
-a `function` object with at least a `name`.
+Ensure your `get_tools()` method returns valid OpenAI-compatible function tool
+schemas. Each tool must have a `type` field set to `"function"` and a `function`
+object with at least a `name`.
### Provider connection errors
`SystemicProviderError` is raised when a provider error affects all rows (e.g.
-authentication failure, budget exhausted, network unreachable). The batch aborts
-early instead of retrying each row.
-
-Resolution -- fix the underlying credential or connectivity issue and re-run.
+authentication failure, budget exhausted). The batch aborts early instead of
+retrying each row. Fix the underlying credential or connectivity issue and re-run.
## Remote Rollout Errors
### AgentLoopValidationError
-Raised by `osmosis validate` or when starting a server with `validate=True`
-(the default). The error lists all validation failures.
-
-```
-AgentLoopValidationError: Agent loop validation failed with 2 error(s):
- - [MISSING_NAME] name: Agent loop must have a 'name' attribute
- - [RUN_NOT_ASYNC] run: 'run' method must be an async function (async def)
-```
+Raised by `osmosis validate` or when starting a server with `validate=True`.
Common validation error codes:
@@ -281,16 +208,11 @@ Common validation error codes:
| `MISSING_NAME` | Agent loop class has no `name` attribute. |
| `INVALID_NAME_TYPE` | `name` is not a string. |
| `EMPTY_NAME` | `name` is empty or whitespace. |
-| `MISSING_RUN_METHOD` | No `run` method defined. |
-| `RUN_NOT_CALLABLE` | `run` exists but is not callable. |
| `RUN_NOT_ASYNC` | `run` is not an `async def` function. |
| `GET_TOOLS_RETURNS_NONE` | `get_tools()` returned `None` instead of a list. |
-| `GET_TOOLS_INVALID_TYPE` | `get_tools()` returned a non-list type. |
| `GET_TOOLS_EXCEPTION` | `get_tools()` raised an exception. |
| `MISSING_TOOL_TYPE` | A tool dict is missing the `type` field. |
-| `MISSING_FUNCTION` | A tool dict is missing the `function` field. |
| `MISSING_FUNCTION_NAME` | A function definition has no `name`. |
-| `INVALID_FUNCTION_NAME` | Function name is empty or not a string. |
Run validation before serving to catch these early:
@@ -300,8 +222,6 @@ osmosis validate -m server:agent_loop
### ServeError
-Raised when `serve_agent_loop()` or `osmosis serve` cannot start the server.
-
**Not logged in:**
```
@@ -311,144 +231,35 @@ ServeError: Not logged in. Please run 'osmosis login' first, or use skip_registe
Resolution -- run `osmosis login`, or pass `--skip-register` / `--local` if you
do not need platform registration.
-**Conflicting options:**
-
-```
-ServeError: local_debug=True disables API key authentication; do not provide api_key in local debug mode.
-```
-
-Resolution -- do not combine `--local` with `--api-key`. Local debug mode
-intentionally disables API key authentication.
-
**Missing dependencies:**
```
ImportError: FastAPI is required for serve_agent_loop(). Install it with: pip install osmosis-ai[server]
```
-```
-ImportError: uvicorn is required for serve_agent_loop(). Install it with: pip install osmosis-ai[server]
-```
-
-Resolution:
-
-```bash
-pip install osmosis-ai[server]
-```
-
### Connection / timeout issues
-Rollout protocol exceptions are defined in
-`osmosis_ai.rollout.core.exceptions`:
+Rollout protocol exceptions:
| Exception | Description | Retryable? |
|-----------|-------------|------------|
| `OsmosisTransportError` | Network-level failure (connection refused, DNS error). | Yes |
-| `OsmosisServerError` | Server returned HTTP 5xx. Includes `status_code` attribute. | Yes |
-| `OsmosisValidationError` | Server returned HTTP 4xx. Includes `status_code` attribute. | No |
+| `OsmosisServerError` | Server returned HTTP 5xx. | Yes |
+| `OsmosisValidationError` | Server returned HTTP 4xx. | No |
| `OsmosisTimeoutError` | Request exceeded configured timeout. | Yes |
-| `AgentLoopNotFoundError` | Registry lookup for agent name failed. Includes `name` and `available` attributes. | No |
-| `ToolExecutionError` | A tool call failed during execution. Includes `tool_call_id` and `tool_name`. | Depends |
-| `ToolArgumentError` | Tool arguments could not be parsed (subclass of `ToolExecutionError`). | No |
+| `AgentLoopNotFoundError` | Registry lookup for agent name failed. | No |
+| `ToolExecutionError` | A tool call failed during execution. | Depends |
+| `ToolArgumentError` | Tool arguments could not be parsed. | No |
-For retryable errors, use exponential backoff. The SDK client settings provide
-configurable retry parameters (see [Configuration](./configuration.md)).
+For retryable errors, use exponential backoff. See [Configuration](./configuration.md) for client retry settings.
### PublicIPDetectionError
-Raised when the server cannot detect its public IP address. This happens when
-binding to `0.0.0.0` and all detection methods fail.
-
-```
-PublicIPDetectionError: Failed to detect public IP address. All detection methods failed:
- 1. Cloud metadata (AWS/GCP/Azure): unavailable or returned no public IP
- 2. External IP services: all failed or timed out
-
-To fix this, provide an explicit IP/hostname to your application.
-```
-
-Resolution -- the SDK tries cloud metadata services (AWS, GCP, Azure) first,
-then external IP services (checkip.amazonaws.com, ipify, icanhazip, ifconfig.me)
-as a fallback. If all fail:
+Raised when the server cannot detect its public IP address. Resolution:
- Check network connectivity and firewall rules.
-- Provide an explicit host via the `OSMOSIS_PUBLIC_HOST` environment variable
- or the `--host` flag.
-- If running locally, use `--local` mode which skips IP detection for
- platform registration.
-
-## Debug Tips
-
-### Using `--verbose` flag
-
-The `osmosis validate` command accepts `-v` / `--verbose` to show detailed
-validation output including warnings.
-
-```bash
-osmosis validate -m server:agent_loop --verbose
-```
-
-### Using `--debug` flag
-
-The `osmosis test` and `osmosis eval` commands accept `--debug` to enable debug
-logging, which prints detailed information about each step including provider
-requests and responses.
-
-```bash
-osmosis test -m server:agent_loop -d data.jsonl --model openai/gpt-5-mini --debug
-osmosis eval -m server:agent_loop -d data.jsonl --eval-fn rewards:fn --model openai/gpt-5-mini --debug
-```
-
-### Interactive mode in test mode
-
-Use `--interactive` with `osmosis test` to step through each row one at a time.
-This is useful for debugging agent logic and inspecting intermediate messages:
-
-```bash
-osmosis test -m server:agent_loop -d data.jsonl --interactive
-
-# Start at a specific row
-osmosis test -m server:agent_loop -d data.jsonl --interactive --row 5
-```
-
-### Debug directory for rollout server
-
-When serving an agent, use `--log` to write detailed execution traces
-for each rollout to JSONL files:
-
-```bash
-osmosis serve -m server:agent_loop --log ./debug-logs
-```
-
-Each rollout will produce a file at `{debug_dir}/{timestamp}/{rollout_id}.jsonl`.
-
-### Checking logs
-
-The SDK uses Python's standard `logging` module. Increase the log level to see
-more detail:
-
-```python
-import logging
-logging.basicConfig(level=logging.DEBUG)
-```
-
-Or set the uvicorn log level when serving:
-
-```bash
-osmosis serve -m server:agent_loop --log-level debug
-```
-
-### Checking credentials
-
-To verify your authentication state and credential file:
-
-```bash
-# Show current user and workspace
-osmosis whoami
-
-# Credentials are stored at:
-# ~/.config/osmosis/credentials.json
-```
+- Provide an explicit host via `OSMOSIS_PUBLIC_HOST` environment variable or the `--host` flag.
+- If running locally, use `--local` mode which skips IP detection.
## See Also
@@ -457,4 +268,3 @@ osmosis whoami
- [Dataset Format](./datasets.md) -- supported formats and required columns
- [Test Mode](./test-mode.md) -- full `osmosis test` documentation
- [Eval Mode](./eval-mode.md) -- full `osmosis eval` documentation
-- [Rewards API Reference](./rewards-api.md) -- `@osmosis_reward`, `@osmosis_rubric`, `evaluate_rubric`
From 95c35df1977ff005af2c1cf502e5f2d9683290c2 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 16:57:11 -0800
Subject: [PATCH 46/60] warning
---
README.md | 2 ++
1 file changed, 2 insertions(+)
diff --git a/README.md b/README.md
index e3d7ed86..c5f1f10b 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,8 @@
# osmosis-ai
+> ⚠️ **Warning**: osmosis-ai is still in active development. APIs may change between versions.
+
Python SDK for Osmosis AI training workflows. Supports two training modes with shared tooling for testing and evaluation.
Osmosis AI is a platform for training LLMs with reinforcement learning. You define custom reward functions, LLM-as-judge rubrics, and agent tools -- then Osmosis handles the training loop on managed GPU clusters. This SDK provides everything you need to build and test those components locally, from `@osmosis_reward` decorators and MCP tool definitions to a full CLI for running agents against datasets before submitting training runs.
From 7123f53ccd13d3ce38f0c0893ffc8591c901abc1 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 17:05:46 -0800
Subject: [PATCH 47/60] improve docstring
---
osmosis_ai/rollout/core/base.py | 6 ++++--
1 file changed, 4 insertions(+), 2 deletions(-)
diff --git a/osmosis_ai/rollout/core/base.py b/osmosis_ai/rollout/core/base.py
index 03b29971..3360ff25 100644
--- a/osmosis_ai/rollout/core/base.py
+++ b/osmosis_ai/rollout/core/base.py
@@ -125,13 +125,15 @@ async def run(self, ctx: RolloutContext) -> RolloutResult:
# Execute tools and add results
for tool_call in result.tool_calls:
args = json.loads(tool_call["function"]["arguments"])
- tool_result = do_tool(tool_call["function"]["name"], args)
+ start = time.monotonic()
+ tool_result = your_tool_fn(tool_call["function"]["name"], args)
+ latency_ms = (time.monotonic() - start) * 1000.0
messages.append({
"role": "tool",
"tool_call_id": tool_call["id"],
"content": str(tool_result),
})
- ctx.record_tool_call(latency_ms=...)
+ ctx.record_tool_call(latency_ms=latency_ms)
return ctx.complete(messages)
"""
From 5a5b61741b6c709645422097a44271e21c1ac550 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Tue, 17 Feb 2026 17:24:09 -0800
Subject: [PATCH 48/60] Remove obsolete build scripts and configuration files;
update .gitignore and GitHub Actions workflow for new build process using
'uv'.
---
.github/workflows/publish.yml | 9 +++++---
.gitignore | 20 ++---------------
MANIFEST.in | 18 ---------------
build_package.bat | 21 ------------------
build_package.sh | 21 ------------------
requirements.lock | 42 -----------------------------------
setup_env.bat | 24 --------------------
7 files changed, 8 insertions(+), 147 deletions(-)
delete mode 100644 MANIFEST.in
delete mode 100644 build_package.bat
delete mode 100755 build_package.sh
delete mode 100644 requirements.lock
delete mode 100644 setup_env.bat
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 0177278c..a153e16e 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -17,13 +17,16 @@ jobs:
- name: Checkout repository
uses: actions/checkout@v6
- - name: Set up Python
- uses: actions/setup-python@v6
+ - name: Set up uv
+ uses: astral-sh/setup-uv@v7
with:
python-version: "3.12"
- name: Build distribution artifacts
- run: ./build_package.sh
+ run: |
+ rm -rf dist/ build/ *.egg-info/
+ uv build
+ uvx twine check --strict dist/*
- name: Publish to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
diff --git a/.gitignore b/.gitignore
index 89691bd8..d4a634ab 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,22 +11,6 @@ __pycache__/
*.py[cod]
*$py.class
*.so
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
# Distribution / packaging
.Python
@@ -63,9 +47,9 @@ coverage.xml
.vscode/
*.swp
*.swo
-*~
+*~
todo.md
.DS_Store
-CLAUDE.md
\ No newline at end of file
+CLAUDE.md
diff --git a/MANIFEST.in b/MANIFEST.in
deleted file mode 100644
index b6a7f24a..00000000
--- a/MANIFEST.in
+++ /dev/null
@@ -1,18 +0,0 @@
-include LICENSE
-include README.md
-include pyproject.toml
-include pytest.ini
-recursive-include osmosis_ai py.typed
-include .env.sample
-include setup_env.sh
-include setup_env.bat
-
-recursive-exclude __pycache__ *
-recursive-exclude *.py[cod] *
-recursive-exclude venv *
-recursive-exclude .pytest_cache *
-recursive-exclude .git *
-recursive-exclude .github *
-recursive-exclude osmosis_ai.egg-info *
-
-prune examples
\ No newline at end of file
diff --git a/build_package.bat b/build_package.bat
deleted file mode 100644
index 58a7dbe7..00000000
--- a/build_package.bat
+++ /dev/null
@@ -1,21 +0,0 @@
-@echo off
-REM Script to build and upload the package to PyPI
-
-REM Clean previous builds
-rmdir /s /q dist build osmosis_ai.egg-info
-
-REM Install build dependencies if not already installed
-pip install --upgrade build twine
-
-REM Build the package
-python -m build
-
-REM Check the package
-twine check dist/*
-
-echo.
-echo To upload to TestPyPI, run:
-echo twine upload --repository testpypi dist/*
-echo.
-echo To upload to PyPI, run:
-echo twine upload dist/*
\ No newline at end of file
diff --git a/build_package.sh b/build_package.sh
deleted file mode 100755
index 2793b4c1..00000000
--- a/build_package.sh
+++ /dev/null
@@ -1,21 +0,0 @@
-#!/bin/bash
-# Script to build and upload the package to PyPI
-
-# Clean previous builds
-rm -rf dist/ build/ *.egg-info/
-
-# Install build dependencies if not already installed
-pip install --upgrade build twine
-
-# Build the package
-python -m build
-
-# Check the package
-twine check dist/*
-
-echo ""
-echo "To upload to TestPyPI, run:"
-echo "twine upload --repository testpypi dist/*"
-echo ""
-echo "To upload to PyPI, run:"
-echo "twine upload dist/*"
\ No newline at end of file
diff --git a/requirements.lock b/requirements.lock
deleted file mode 100644
index a91c0ae4..00000000
--- a/requirements.lock
+++ /dev/null
@@ -1,42 +0,0 @@
-annotated-types==0.7.0
-anthropic==0.49.0
-anyio==4.9.0
-black==25.1.0
-certifi==2025.1.31
-charset-normalizer==3.4.1
-click==8.1.8
-distro==1.9.0
-h11==0.14.0
-httpcore==1.0.7
-httpx==0.28.1
-idna==3.10
-iniconfig==2.1.0
-isort==6.0.1
-jiter==0.9.0
-jsonpatch==1.33
-jsonpointer==3.0.0
-langsmith==0.3.18
-mypy-extensions==1.0.0
-openai==2.0.0
-orjson==3.10.16
-packaging==24.2
-pathspec==0.12.1
-platformdirs==4.3.7
-pluggy==1.5.0
-pydantic==2.10.6
-pydantic_core==2.27.2
-pytest==8.3.5
-python-dotenv==1.0.1
-PyYAML==6.0.2
-regex==2024.11.6
-requests==2.32.3
-requests-toolbelt==1.0.0
-sniffio==1.3.1
-SQLAlchemy==2.0.39
-tenacity==9.0.0
-tiktoken==0.9.0
-tqdm==4.67.1
-typing_extensions==4.12.2
-urllib3==2.3.0
-xxhash==3.5.0
-zstandard==0.23.0
diff --git a/setup_env.bat b/setup_env.bat
deleted file mode 100644
index c5baff59..00000000
--- a/setup_env.bat
+++ /dev/null
@@ -1,24 +0,0 @@
-@echo off
-REM Script to set up a Python virtual environment for osmosis-ai on Windows
-
-echo Creating virtual environment...
-python -m venv venv
-
-echo Activating virtual environment...
-call venv\Scripts\activate.bat
-
-echo Installing dependencies...
-pip install -e ".[dev]"
-
-if not exist .env (
- echo Creating .env file from template...
- copy .env.sample .env
- echo Please edit .env file to add your API keys
-)
-
-echo.
-echo Environment setup complete!
-echo To activate the environment, run: venv\Scripts\activate.bat
-echo Remember to edit your .env file to add your API keys.
-
-pause
\ No newline at end of file
From 3b7561864ebea6e3e3bb4eb3f752ed2a8727aba5 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 09:16:26 -0800
Subject: [PATCH 49/60] address Dependabot alerts
---
uv.lock | 216 ++++++++++++++++++++++++++++----------------------------
1 file changed, 108 insertions(+), 108 deletions(-)
diff --git a/uv.lock b/uv.lock
index d86de8d2..c931705c 100644
--- a/uv.lock
+++ b/uv.lock
@@ -17,7 +17,7 @@ wheels = [
[[package]]
name = "aiohttp"
-version = "3.13.2"
+version = "3.13.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohappyeyeballs" },
@@ -29,110 +29,110 @@ dependencies = [
{ name = "propcache" },
{ name = "yarl" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/1c/ce/3b83ebba6b3207a7135e5fcaba49706f8a4b6008153b4e30540c982fae26/aiohttp-3.13.2.tar.gz", hash = "sha256:40176a52c186aefef6eb3cad2cdd30cd06e3afbe88fe8ab2af9c0b90f228daca", size = 7837994, upload-time = "2025-10-28T20:59:39.937Z" }
-wheels = [
- { url = "https://files.pythonhosted.org/packages/6d/34/939730e66b716b76046dedfe0842995842fa906ccc4964bba414ff69e429/aiohttp-3.13.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:2372b15a5f62ed37789a6b383ff7344fc5b9f243999b0cd9b629d8bc5f5b4155", size = 736471, upload-time = "2025-10-28T20:55:27.924Z" },
- { url = "https://files.pythonhosted.org/packages/fd/cf/dcbdf2df7f6ca72b0bb4c0b4509701f2d8942cf54e29ca197389c214c07f/aiohttp-3.13.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7f8659a48995edee7229522984bd1009c1213929c769c2daa80b40fe49a180c", size = 493985, upload-time = "2025-10-28T20:55:29.456Z" },
- { url = "https://files.pythonhosted.org/packages/9d/87/71c8867e0a1d0882dcbc94af767784c3cb381c1c4db0943ab4aae4fed65e/aiohttp-3.13.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:939ced4a7add92296b0ad38892ce62b98c619288a081170695c6babe4f50e636", size = 489274, upload-time = "2025-10-28T20:55:31.134Z" },
- { url = "https://files.pythonhosted.org/packages/38/0f/46c24e8dae237295eaadd113edd56dee96ef6462adf19b88592d44891dc5/aiohttp-3.13.2-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6315fb6977f1d0dd41a107c527fee2ed5ab0550b7d885bc15fee20ccb17891da", size = 1668171, upload-time = "2025-10-28T20:55:36.065Z" },
- { url = "https://files.pythonhosted.org/packages/eb/c6/4cdfb4440d0e28483681a48f69841fa5e39366347d66ef808cbdadddb20e/aiohttp-3.13.2-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6e7352512f763f760baaed2637055c49134fd1d35b37c2dedfac35bfe5cf8725", size = 1636036, upload-time = "2025-10-28T20:55:37.576Z" },
- { url = "https://files.pythonhosted.org/packages/84/37/8708cf678628216fb678ab327a4e1711c576d6673998f4f43e86e9ae90dd/aiohttp-3.13.2-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e09a0a06348a2dd73e7213353c90d709502d9786219f69b731f6caa0efeb46f5", size = 1727975, upload-time = "2025-10-28T20:55:39.457Z" },
- { url = "https://files.pythonhosted.org/packages/e6/2e/3ebfe12fdcb9b5f66e8a0a42dffcd7636844c8a018f261efb2419f68220b/aiohttp-3.13.2-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a09a6d073fb5789456545bdee2474d14395792faa0527887f2f4ec1a486a59d3", size = 1815823, upload-time = "2025-10-28T20:55:40.958Z" },
- { url = "https://files.pythonhosted.org/packages/a1/4f/ca2ef819488cbb41844c6cf92ca6dd15b9441e6207c58e5ae0e0fc8d70ad/aiohttp-3.13.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b59d13c443f8e049d9e94099c7e412e34610f1f49be0f230ec656a10692a5802", size = 1669374, upload-time = "2025-10-28T20:55:42.745Z" },
- { url = "https://files.pythonhosted.org/packages/f8/fe/1fe2e1179a0d91ce09c99069684aab619bf2ccde9b20bd6ca44f8837203e/aiohttp-3.13.2-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:20db2d67985d71ca033443a1ba2001c4b5693fe09b0e29f6d9358a99d4d62a8a", size = 1555315, upload-time = "2025-10-28T20:55:44.264Z" },
- { url = "https://files.pythonhosted.org/packages/5a/2b/f3781899b81c45d7cbc7140cddb8a3481c195e7cbff8e36374759d2ab5a5/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:960c2fc686ba27b535f9fd2b52d87ecd7e4fd1cf877f6a5cba8afb5b4a8bd204", size = 1639140, upload-time = "2025-10-28T20:55:46.626Z" },
- { url = "https://files.pythonhosted.org/packages/72/27/c37e85cd3ece6f6c772e549bd5a253d0c122557b25855fb274224811e4f2/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:6c00dbcf5f0d88796151e264a8eab23de2997c9303dd7c0bf622e23b24d3ce22", size = 1645496, upload-time = "2025-10-28T20:55:48.933Z" },
- { url = "https://files.pythonhosted.org/packages/66/20/3af1ab663151bd3780b123e907761cdb86ec2c4e44b2d9b195ebc91fbe37/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:fed38a5edb7945f4d1bcabe2fcd05db4f6ec7e0e82560088b754f7e08d93772d", size = 1697625, upload-time = "2025-10-28T20:55:50.377Z" },
- { url = "https://files.pythonhosted.org/packages/95/eb/ae5cab15efa365e13d56b31b0d085a62600298bf398a7986f8388f73b598/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:b395bbca716c38bef3c764f187860e88c724b342c26275bc03e906142fc5964f", size = 1542025, upload-time = "2025-10-28T20:55:51.861Z" },
- { url = "https://files.pythonhosted.org/packages/e9/2d/1683e8d67ec72d911397fe4e575688d2a9b8f6a6e03c8fdc9f3fd3d4c03f/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:204ffff2426c25dfda401ba08da85f9c59525cdc42bda26660463dd1cbcfec6f", size = 1714918, upload-time = "2025-10-28T20:55:53.515Z" },
- { url = "https://files.pythonhosted.org/packages/99/a2/ffe8e0e1c57c5e542d47ffa1fcf95ef2b3ea573bf7c4d2ee877252431efc/aiohttp-3.13.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:05c4dd3c48fb5f15db31f57eb35374cb0c09afdde532e7fb70a75aede0ed30f6", size = 1656113, upload-time = "2025-10-28T20:55:55.438Z" },
- { url = "https://files.pythonhosted.org/packages/0d/42/d511aff5c3a2b06c09d7d214f508a4ad8ac7799817f7c3d23e7336b5e896/aiohttp-3.13.2-cp310-cp310-win32.whl", hash = "sha256:e574a7d61cf10351d734bcddabbe15ede0eaa8a02070d85446875dc11189a251", size = 432290, upload-time = "2025-10-28T20:55:56.96Z" },
- { url = "https://files.pythonhosted.org/packages/8b/ea/1c2eb7098b5bad4532994f2b7a8228d27674035c9b3234fe02c37469ef14/aiohttp-3.13.2-cp310-cp310-win_amd64.whl", hash = "sha256:364f55663085d658b8462a1c3f17b2b84a5c2e1ba858e1b79bff7b2e24ad1514", size = 455075, upload-time = "2025-10-28T20:55:58.373Z" },
- { url = "https://files.pythonhosted.org/packages/35/74/b321e7d7ca762638cdf8cdeceb39755d9c745aff7a64c8789be96ddf6e96/aiohttp-3.13.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:4647d02df098f6434bafd7f32ad14942f05a9caa06c7016fdcc816f343997dd0", size = 743409, upload-time = "2025-10-28T20:56:00.354Z" },
- { url = "https://files.pythonhosted.org/packages/99/3d/91524b905ec473beaf35158d17f82ef5a38033e5809fe8742e3657cdbb97/aiohttp-3.13.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:e3403f24bcb9c3b29113611c3c16a2a447c3953ecf86b79775e7be06f7ae7ccb", size = 497006, upload-time = "2025-10-28T20:56:01.85Z" },
- { url = "https://files.pythonhosted.org/packages/eb/d3/7f68bc02a67716fe80f063e19adbd80a642e30682ce74071269e17d2dba1/aiohttp-3.13.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:43dff14e35aba17e3d6d5ba628858fb8cb51e30f44724a2d2f0c75be492c55e9", size = 493195, upload-time = "2025-10-28T20:56:03.314Z" },
- { url = "https://files.pythonhosted.org/packages/98/31/913f774a4708775433b7375c4f867d58ba58ead833af96c8af3621a0d243/aiohttp-3.13.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e2a9ea08e8c58bb17655630198833109227dea914cd20be660f52215f6de5613", size = 1747759, upload-time = "2025-10-28T20:56:04.904Z" },
- { url = "https://files.pythonhosted.org/packages/e8/63/04efe156f4326f31c7c4a97144f82132c3bb21859b7bb84748d452ccc17c/aiohttp-3.13.2-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:53b07472f235eb80e826ad038c9d106c2f653584753f3ddab907c83f49eedead", size = 1704456, upload-time = "2025-10-28T20:56:06.986Z" },
- { url = "https://files.pythonhosted.org/packages/8e/02/4e16154d8e0a9cf4ae76f692941fd52543bbb148f02f098ca73cab9b1c1b/aiohttp-3.13.2-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e736c93e9c274fce6419af4aac199984d866e55f8a4cec9114671d0ea9688780", size = 1807572, upload-time = "2025-10-28T20:56:08.558Z" },
- { url = "https://files.pythonhosted.org/packages/34/58/b0583defb38689e7f06798f0285b1ffb3a6fb371f38363ce5fd772112724/aiohttp-3.13.2-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ff5e771f5dcbc81c64898c597a434f7682f2259e0cd666932a913d53d1341d1a", size = 1895954, upload-time = "2025-10-28T20:56:10.545Z" },
- { url = "https://files.pythonhosted.org/packages/6b/f3/083907ee3437425b4e376aa58b2c915eb1a33703ec0dc30040f7ae3368c6/aiohttp-3.13.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3b6fb0c207cc661fa0bf8c66d8d9b657331ccc814f4719468af61034b478592", size = 1747092, upload-time = "2025-10-28T20:56:12.118Z" },
- { url = "https://files.pythonhosted.org/packages/ac/61/98a47319b4e425cc134e05e5f3fc512bf9a04bf65aafd9fdcda5d57ec693/aiohttp-3.13.2-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:97a0895a8e840ab3520e2288db7cace3a1981300d48babeb50e7425609e2e0ab", size = 1606815, upload-time = "2025-10-28T20:56:14.191Z" },
- { url = "https://files.pythonhosted.org/packages/97/4b/e78b854d82f66bb974189135d31fce265dee0f5344f64dd0d345158a5973/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9e8f8afb552297aca127c90cb840e9a1d4bfd6a10d7d8f2d9176e1acc69bad30", size = 1723789, upload-time = "2025-10-28T20:56:16.101Z" },
- { url = "https://files.pythonhosted.org/packages/ed/fc/9d2ccc794fc9b9acd1379d625c3a8c64a45508b5091c546dea273a41929e/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:ed2f9c7216e53c3df02264f25d824b079cc5914f9e2deba94155190ef648ee40", size = 1718104, upload-time = "2025-10-28T20:56:17.655Z" },
- { url = "https://files.pythonhosted.org/packages/66/65/34564b8765ea5c7d79d23c9113135d1dd3609173da13084830f1507d56cf/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:99c5280a329d5fa18ef30fd10c793a190d996567667908bef8a7f81f8202b948", size = 1785584, upload-time = "2025-10-28T20:56:19.238Z" },
- { url = "https://files.pythonhosted.org/packages/30/be/f6a7a426e02fc82781afd62016417b3948e2207426d90a0e478790d1c8a4/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:2ca6ffef405fc9c09a746cb5d019c1672cd7f402542e379afc66b370833170cf", size = 1595126, upload-time = "2025-10-28T20:56:20.836Z" },
- { url = "https://files.pythonhosted.org/packages/e5/c7/8e22d5d28f94f67d2af496f14a83b3c155d915d1fe53d94b66d425ec5b42/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:47f438b1a28e926c37632bff3c44df7d27c9b57aaf4e34b1def3c07111fdb782", size = 1800665, upload-time = "2025-10-28T20:56:22.922Z" },
- { url = "https://files.pythonhosted.org/packages/d1/11/91133c8b68b1da9fc16555706aa7276fdf781ae2bb0876c838dd86b8116e/aiohttp-3.13.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9acda8604a57bb60544e4646a4615c1866ee6c04a8edef9b8ee6fd1d8fa2ddc8", size = 1739532, upload-time = "2025-10-28T20:56:25.924Z" },
- { url = "https://files.pythonhosted.org/packages/17/6b/3747644d26a998774b21a616016620293ddefa4d63af6286f389aedac844/aiohttp-3.13.2-cp311-cp311-win32.whl", hash = "sha256:868e195e39b24aaa930b063c08bb0c17924899c16c672a28a65afded9c46c6ec", size = 431876, upload-time = "2025-10-28T20:56:27.524Z" },
- { url = "https://files.pythonhosted.org/packages/c3/63/688462108c1a00eb9f05765331c107f95ae86f6b197b865d29e930b7e462/aiohttp-3.13.2-cp311-cp311-win_amd64.whl", hash = "sha256:7fd19df530c292542636c2a9a85854fab93474396a52f1695e799186bbd7f24c", size = 456205, upload-time = "2025-10-28T20:56:29.062Z" },
- { url = "https://files.pythonhosted.org/packages/29/9b/01f00e9856d0a73260e86dd8ed0c2234a466c5c1712ce1c281548df39777/aiohttp-3.13.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:b1e56bab2e12b2b9ed300218c351ee2a3d8c8fdab5b1ec6193e11a817767e47b", size = 737623, upload-time = "2025-10-28T20:56:30.797Z" },
- { url = "https://files.pythonhosted.org/packages/5a/1b/4be39c445e2b2bd0aab4ba736deb649fabf14f6757f405f0c9685019b9e9/aiohttp-3.13.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:364e25edaabd3d37b1db1f0cbcee8c73c9a3727bfa262b83e5e4cf3489a2a9dc", size = 492664, upload-time = "2025-10-28T20:56:32.708Z" },
- { url = "https://files.pythonhosted.org/packages/28/66/d35dcfea8050e131cdd731dff36434390479b4045a8d0b9d7111b0a968f1/aiohttp-3.13.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c5c94825f744694c4b8db20b71dba9a257cd2ba8e010a803042123f3a25d50d7", size = 491808, upload-time = "2025-10-28T20:56:34.57Z" },
- { url = "https://files.pythonhosted.org/packages/00/29/8e4609b93e10a853b65f8291e64985de66d4f5848c5637cddc70e98f01f8/aiohttp-3.13.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ba2715d842ffa787be87cbfce150d5e88c87a98e0b62e0f5aa489169a393dbbb", size = 1738863, upload-time = "2025-10-28T20:56:36.377Z" },
- { url = "https://files.pythonhosted.org/packages/9d/fa/4ebdf4adcc0def75ced1a0d2d227577cd7b1b85beb7edad85fcc87693c75/aiohttp-3.13.2-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:585542825c4bc662221fb257889e011a5aa00f1ae4d75d1d246a5225289183e3", size = 1700586, upload-time = "2025-10-28T20:56:38.034Z" },
- { url = "https://files.pythonhosted.org/packages/da/04/73f5f02ff348a3558763ff6abe99c223381b0bace05cd4530a0258e52597/aiohttp-3.13.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:39d02cb6025fe1aabca329c5632f48c9532a3dabccd859e7e2f110668972331f", size = 1768625, upload-time = "2025-10-28T20:56:39.75Z" },
- { url = "https://files.pythonhosted.org/packages/f8/49/a825b79ffec124317265ca7d2344a86bcffeb960743487cb11988ffb3494/aiohttp-3.13.2-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e67446b19e014d37342f7195f592a2a948141d15a312fe0e700c2fd2f03124f6", size = 1867281, upload-time = "2025-10-28T20:56:41.471Z" },
- { url = "https://files.pythonhosted.org/packages/b9/48/adf56e05f81eac31edcfae45c90928f4ad50ef2e3ea72cb8376162a368f8/aiohttp-3.13.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4356474ad6333e41ccefd39eae869ba15a6c5299c9c01dfdcfdd5c107be4363e", size = 1752431, upload-time = "2025-10-28T20:56:43.162Z" },
- { url = "https://files.pythonhosted.org/packages/30/ab/593855356eead019a74e862f21523db09c27f12fd24af72dbc3555b9bfd9/aiohttp-3.13.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:eeacf451c99b4525f700f078becff32c32ec327b10dcf31306a8a52d78166de7", size = 1562846, upload-time = "2025-10-28T20:56:44.85Z" },
- { url = "https://files.pythonhosted.org/packages/39/0f/9f3d32271aa8dc35036e9668e31870a9d3b9542dd6b3e2c8a30931cb27ae/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d8a9b889aeabd7a4e9af0b7f4ab5ad94d42e7ff679aaec6d0db21e3b639ad58d", size = 1699606, upload-time = "2025-10-28T20:56:46.519Z" },
- { url = "https://files.pythonhosted.org/packages/2c/3c/52d2658c5699b6ef7692a3f7128b2d2d4d9775f2a68093f74bca06cf01e1/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:fa89cb11bc71a63b69568d5b8a25c3ca25b6d54c15f907ca1c130d72f320b76b", size = 1720663, upload-time = "2025-10-28T20:56:48.528Z" },
- { url = "https://files.pythonhosted.org/packages/9b/d4/8f8f3ff1fb7fb9e3f04fcad4e89d8a1cd8fc7d05de67e3de5b15b33008ff/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8aa7c807df234f693fed0ecd507192fc97692e61fee5702cdc11155d2e5cadc8", size = 1737939, upload-time = "2025-10-28T20:56:50.77Z" },
- { url = "https://files.pythonhosted.org/packages/03/d3/ddd348f8a27a634daae39a1b8e291ff19c77867af438af844bf8b7e3231b/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:9eb3e33fdbe43f88c3c75fa608c25e7c47bbd80f48d012763cb67c47f39a7e16", size = 1555132, upload-time = "2025-10-28T20:56:52.568Z" },
- { url = "https://files.pythonhosted.org/packages/39/b8/46790692dc46218406f94374903ba47552f2f9f90dad554eed61bfb7b64c/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:9434bc0d80076138ea986833156c5a48c9c7a8abb0c96039ddbb4afc93184169", size = 1764802, upload-time = "2025-10-28T20:56:54.292Z" },
- { url = "https://files.pythonhosted.org/packages/ba/e4/19ce547b58ab2a385e5f0b8aa3db38674785085abcf79b6e0edd1632b12f/aiohttp-3.13.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ff15c147b2ad66da1f2cbb0622313f2242d8e6e8f9b79b5206c84523a4473248", size = 1719512, upload-time = "2025-10-28T20:56:56.428Z" },
- { url = "https://files.pythonhosted.org/packages/70/30/6355a737fed29dcb6dfdd48682d5790cb5eab050f7b4e01f49b121d3acad/aiohttp-3.13.2-cp312-cp312-win32.whl", hash = "sha256:27e569eb9d9e95dbd55c0fc3ec3a9335defbf1d8bc1d20171a49f3c4c607b93e", size = 426690, upload-time = "2025-10-28T20:56:58.736Z" },
- { url = "https://files.pythonhosted.org/packages/0a/0d/b10ac09069973d112de6ef980c1f6bb31cb7dcd0bc363acbdad58f927873/aiohttp-3.13.2-cp312-cp312-win_amd64.whl", hash = "sha256:8709a0f05d59a71f33fd05c17fc11fcb8c30140506e13c2f5e8ee1b8964e1b45", size = 453465, upload-time = "2025-10-28T20:57:00.795Z" },
- { url = "https://files.pythonhosted.org/packages/bf/78/7e90ca79e5aa39f9694dcfd74f4720782d3c6828113bb1f3197f7e7c4a56/aiohttp-3.13.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:7519bdc7dfc1940d201651b52bf5e03f5503bda45ad6eacf64dda98be5b2b6be", size = 732139, upload-time = "2025-10-28T20:57:02.455Z" },
- { url = "https://files.pythonhosted.org/packages/db/ed/1f59215ab6853fbaa5c8495fa6cbc39edfc93553426152b75d82a5f32b76/aiohttp-3.13.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:088912a78b4d4f547a1f19c099d5a506df17eacec3c6f4375e2831ec1d995742", size = 490082, upload-time = "2025-10-28T20:57:04.784Z" },
- { url = "https://files.pythonhosted.org/packages/68/7b/fe0fe0f5e05e13629d893c760465173a15ad0039c0a5b0d0040995c8075e/aiohttp-3.13.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5276807b9de9092af38ed23ce120539ab0ac955547b38563a9ba4f5b07b95293", size = 489035, upload-time = "2025-10-28T20:57:06.894Z" },
- { url = "https://files.pythonhosted.org/packages/d2/04/db5279e38471b7ac801d7d36a57d1230feeee130bbe2a74f72731b23c2b1/aiohttp-3.13.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1237c1375eaef0db4dcd7c2559f42e8af7b87ea7d295b118c60c36a6e61cb811", size = 1720387, upload-time = "2025-10-28T20:57:08.685Z" },
- { url = "https://files.pythonhosted.org/packages/31/07/8ea4326bd7dae2bd59828f69d7fdc6e04523caa55e4a70f4a8725a7e4ed2/aiohttp-3.13.2-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:96581619c57419c3d7d78703d5b78c1e5e5fc0172d60f555bdebaced82ded19a", size = 1688314, upload-time = "2025-10-28T20:57:10.693Z" },
- { url = "https://files.pythonhosted.org/packages/48/ab/3d98007b5b87ffd519d065225438cc3b668b2f245572a8cb53da5dd2b1bc/aiohttp-3.13.2-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a2713a95b47374169409d18103366de1050fe0ea73db358fc7a7acb2880422d4", size = 1756317, upload-time = "2025-10-28T20:57:12.563Z" },
- { url = "https://files.pythonhosted.org/packages/97/3d/801ca172b3d857fafb7b50c7c03f91b72b867a13abca982ed6b3081774ef/aiohttp-3.13.2-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:228a1cd556b3caca590e9511a89444925da87d35219a49ab5da0c36d2d943a6a", size = 1858539, upload-time = "2025-10-28T20:57:14.623Z" },
- { url = "https://files.pythonhosted.org/packages/f7/0d/4764669bdf47bd472899b3d3db91fffbe925c8e3038ec591a2fd2ad6a14d/aiohttp-3.13.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ac6cde5fba8d7d8c6ac963dbb0256a9854e9fafff52fbcc58fdf819357892c3e", size = 1739597, upload-time = "2025-10-28T20:57:16.399Z" },
- { url = "https://files.pythonhosted.org/packages/c4/52/7bd3c6693da58ba16e657eb904a5b6decfc48ecd06e9ac098591653b1566/aiohttp-3.13.2-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f2bef8237544f4e42878c61cef4e2839fee6346dc60f5739f876a9c50be7fcdb", size = 1555006, upload-time = "2025-10-28T20:57:18.288Z" },
- { url = "https://files.pythonhosted.org/packages/48/30/9586667acec5993b6f41d2ebcf96e97a1255a85f62f3c653110a5de4d346/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:16f15a4eac3bc2d76c45f7ebdd48a65d41b242eb6c31c2245463b40b34584ded", size = 1683220, upload-time = "2025-10-28T20:57:20.241Z" },
- { url = "https://files.pythonhosted.org/packages/71/01/3afe4c96854cfd7b30d78333852e8e851dceaec1c40fd00fec90c6402dd2/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:bb7fb776645af5cc58ab804c58d7eba545a97e047254a52ce89c157b5af6cd0b", size = 1712570, upload-time = "2025-10-28T20:57:22.253Z" },
- { url = "https://files.pythonhosted.org/packages/11/2c/22799d8e720f4697a9e66fd9c02479e40a49de3de2f0bbe7f9f78a987808/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:e1b4951125ec10c70802f2cb09736c895861cd39fd9dcb35107b4dc8ae6220b8", size = 1733407, upload-time = "2025-10-28T20:57:24.37Z" },
- { url = "https://files.pythonhosted.org/packages/34/cb/90f15dd029f07cebbd91f8238a8b363978b530cd128488085b5703683594/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:550bf765101ae721ee1d37d8095f47b1f220650f85fe1af37a90ce75bab89d04", size = 1550093, upload-time = "2025-10-28T20:57:26.257Z" },
- { url = "https://files.pythonhosted.org/packages/69/46/12dce9be9d3303ecbf4d30ad45a7683dc63d90733c2d9fe512be6716cd40/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:fe91b87fc295973096251e2d25a811388e7d8adf3bd2b97ef6ae78bc4ac6c476", size = 1758084, upload-time = "2025-10-28T20:57:28.349Z" },
- { url = "https://files.pythonhosted.org/packages/f9/c8/0932b558da0c302ffd639fc6362a313b98fdf235dc417bc2493da8394df7/aiohttp-3.13.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e0c8e31cfcc4592cb200160344b2fb6ae0f9e4effe06c644b5a125d4ae5ebe23", size = 1716987, upload-time = "2025-10-28T20:57:30.233Z" },
- { url = "https://files.pythonhosted.org/packages/5d/8b/f5bd1a75003daed099baec373aed678f2e9b34f2ad40d85baa1368556396/aiohttp-3.13.2-cp313-cp313-win32.whl", hash = "sha256:0740f31a60848d6edb296a0df827473eede90c689b8f9f2a4cdde74889eb2254", size = 425859, upload-time = "2025-10-28T20:57:32.105Z" },
- { url = "https://files.pythonhosted.org/packages/5d/28/a8a9fc6957b2cee8902414e41816b5ab5536ecf43c3b1843c10e82c559b2/aiohttp-3.13.2-cp313-cp313-win_amd64.whl", hash = "sha256:a88d13e7ca367394908f8a276b89d04a3652044612b9a408a0bb22a5ed976a1a", size = 452192, upload-time = "2025-10-28T20:57:34.166Z" },
- { url = "https://files.pythonhosted.org/packages/9b/36/e2abae1bd815f01c957cbf7be817b3043304e1c87bad526292a0410fdcf9/aiohttp-3.13.2-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:2475391c29230e063ef53a66669b7b691c9bfc3f1426a0f7bcdf1216bdbac38b", size = 735234, upload-time = "2025-10-28T20:57:36.415Z" },
- { url = "https://files.pythonhosted.org/packages/ca/e3/1ee62dde9b335e4ed41db6bba02613295a0d5b41f74a783c142745a12763/aiohttp-3.13.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:f33c8748abef4d8717bb20e8fb1b3e07c6adacb7fd6beaae971a764cf5f30d61", size = 490733, upload-time = "2025-10-28T20:57:38.205Z" },
- { url = "https://files.pythonhosted.org/packages/1a/aa/7a451b1d6a04e8d15a362af3e9b897de71d86feac3babf8894545d08d537/aiohttp-3.13.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ae32f24bbfb7dbb485a24b30b1149e2f200be94777232aeadba3eecece4d0aa4", size = 491303, upload-time = "2025-10-28T20:57:40.122Z" },
- { url = "https://files.pythonhosted.org/packages/57/1e/209958dbb9b01174870f6a7538cd1f3f28274fdbc88a750c238e2c456295/aiohttp-3.13.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5d7f02042c1f009ffb70067326ef183a047425bb2ff3bc434ead4dd4a4a66a2b", size = 1717965, upload-time = "2025-10-28T20:57:42.28Z" },
- { url = "https://files.pythonhosted.org/packages/08/aa/6a01848d6432f241416bc4866cae8dc03f05a5a884d2311280f6a09c73d6/aiohttp-3.13.2-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:93655083005d71cd6c072cdab54c886e6570ad2c4592139c3fb967bfc19e4694", size = 1667221, upload-time = "2025-10-28T20:57:44.869Z" },
- { url = "https://files.pythonhosted.org/packages/87/4f/36c1992432d31bbc789fa0b93c768d2e9047ec8c7177e5cd84ea85155f36/aiohttp-3.13.2-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:0db1e24b852f5f664cd728db140cf11ea0e82450471232a394b3d1a540b0f906", size = 1757178, upload-time = "2025-10-28T20:57:47.216Z" },
- { url = "https://files.pythonhosted.org/packages/ac/b4/8e940dfb03b7e0f68a82b88fd182b9be0a65cb3f35612fe38c038c3112cf/aiohttp-3.13.2-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b009194665bcd128e23eaddef362e745601afa4641930848af4c8559e88f18f9", size = 1838001, upload-time = "2025-10-28T20:57:49.337Z" },
- { url = "https://files.pythonhosted.org/packages/d7/ef/39f3448795499c440ab66084a9db7d20ca7662e94305f175a80f5b7e0072/aiohttp-3.13.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c038a8fdc8103cd51dbd986ecdce141473ffd9775a7a8057a6ed9c3653478011", size = 1716325, upload-time = "2025-10-28T20:57:51.327Z" },
- { url = "https://files.pythonhosted.org/packages/d7/51/b311500ffc860b181c05d91c59a1313bdd05c82960fdd4035a15740d431e/aiohttp-3.13.2-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:66bac29b95a00db411cd758fea0e4b9bdba6d549dfe333f9a945430f5f2cc5a6", size = 1547978, upload-time = "2025-10-28T20:57:53.554Z" },
- { url = "https://files.pythonhosted.org/packages/31/64/b9d733296ef79815226dab8c586ff9e3df41c6aff2e16c06697b2d2e6775/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:4ebf9cfc9ba24a74cf0718f04aac2a3bbe745902cc7c5ebc55c0f3b5777ef213", size = 1682042, upload-time = "2025-10-28T20:57:55.617Z" },
- { url = "https://files.pythonhosted.org/packages/3f/30/43d3e0f9d6473a6db7d472104c4eff4417b1e9df01774cb930338806d36b/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:a4b88ebe35ce54205c7074f7302bd08a4cb83256a3e0870c72d6f68a3aaf8e49", size = 1680085, upload-time = "2025-10-28T20:57:57.59Z" },
- { url = "https://files.pythonhosted.org/packages/16/51/c709f352c911b1864cfd1087577760ced64b3e5bee2aa88b8c0c8e2e4972/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:98c4fb90bb82b70a4ed79ca35f656f4281885be076f3f970ce315402b53099ae", size = 1728238, upload-time = "2025-10-28T20:57:59.525Z" },
- { url = "https://files.pythonhosted.org/packages/19/e2/19bd4c547092b773caeb48ff5ae4b1ae86756a0ee76c16727fcfd281404b/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:ec7534e63ae0f3759df3a1ed4fa6bc8f75082a924b590619c0dd2f76d7043caa", size = 1544395, upload-time = "2025-10-28T20:58:01.914Z" },
- { url = "https://files.pythonhosted.org/packages/cf/87/860f2803b27dfc5ed7be532832a3498e4919da61299b4a1f8eb89b8ff44d/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:5b927cf9b935a13e33644cbed6c8c4b2d0f25b713d838743f8fe7191b33829c4", size = 1742965, upload-time = "2025-10-28T20:58:03.972Z" },
- { url = "https://files.pythonhosted.org/packages/67/7f/db2fc7618925e8c7a601094d5cbe539f732df4fb570740be88ed9e40e99a/aiohttp-3.13.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:88d6c017966a78c5265d996c19cdb79235be5e6412268d7e2ce7dee339471b7a", size = 1697585, upload-time = "2025-10-28T20:58:06.189Z" },
- { url = "https://files.pythonhosted.org/packages/0c/07/9127916cb09bb38284db5036036042b7b2c514c8ebaeee79da550c43a6d6/aiohttp-3.13.2-cp314-cp314-win32.whl", hash = "sha256:f7c183e786e299b5d6c49fb43a769f8eb8e04a2726a2bd5887b98b5cc2d67940", size = 431621, upload-time = "2025-10-28T20:58:08.636Z" },
- { url = "https://files.pythonhosted.org/packages/fb/41/554a8a380df6d3a2bba8a7726429a23f4ac62aaf38de43bb6d6cde7b4d4d/aiohttp-3.13.2-cp314-cp314-win_amd64.whl", hash = "sha256:fe242cd381e0fb65758faf5ad96c2e460df6ee5b2de1072fe97e4127927e00b4", size = 457627, upload-time = "2025-10-28T20:58:11Z" },
- { url = "https://files.pythonhosted.org/packages/c7/8e/3824ef98c039d3951cb65b9205a96dd2b20f22241ee17d89c5701557c826/aiohttp-3.13.2-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:f10d9c0b0188fe85398c61147bbd2a657d616c876863bfeff43376e0e3134673", size = 767360, upload-time = "2025-10-28T20:58:13.358Z" },
- { url = "https://files.pythonhosted.org/packages/a4/0f/6a03e3fc7595421274fa34122c973bde2d89344f8a881b728fa8c774e4f1/aiohttp-3.13.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e7c952aefdf2460f4ae55c5e9c3e80aa72f706a6317e06020f80e96253b1accd", size = 504616, upload-time = "2025-10-28T20:58:15.339Z" },
- { url = "https://files.pythonhosted.org/packages/c6/aa/ed341b670f1bc8a6f2c6a718353d13b9546e2cef3544f573c6a1ff0da711/aiohttp-3.13.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c20423ce14771d98353d2e25e83591fa75dfa90a3c1848f3d7c68243b4fbded3", size = 509131, upload-time = "2025-10-28T20:58:17.693Z" },
- { url = "https://files.pythonhosted.org/packages/7f/f0/c68dac234189dae5c4bbccc0f96ce0cc16b76632cfc3a08fff180045cfa4/aiohttp-3.13.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e96eb1a34396e9430c19d8338d2ec33015e4a87ef2b4449db94c22412e25ccdf", size = 1864168, upload-time = "2025-10-28T20:58:20.113Z" },
- { url = "https://files.pythonhosted.org/packages/8f/65/75a9a76db8364b5d0e52a0c20eabc5d52297385d9af9c35335b924fafdee/aiohttp-3.13.2-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:23fb0783bc1a33640036465019d3bba069942616a6a2353c6907d7fe1ccdaf4e", size = 1719200, upload-time = "2025-10-28T20:58:22.583Z" },
- { url = "https://files.pythonhosted.org/packages/f5/55/8df2ed78d7f41d232f6bd3ff866b6f617026551aa1d07e2f03458f964575/aiohttp-3.13.2-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2e1a9bea6244a1d05a4e57c295d69e159a5c50d8ef16aa390948ee873478d9a5", size = 1843497, upload-time = "2025-10-28T20:58:24.672Z" },
- { url = "https://files.pythonhosted.org/packages/e9/e0/94d7215e405c5a02ccb6a35c7a3a6cfff242f457a00196496935f700cde5/aiohttp-3.13.2-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0a3d54e822688b56e9f6b5816fb3de3a3a64660efac64e4c2dc435230ad23bad", size = 1935703, upload-time = "2025-10-28T20:58:26.758Z" },
- { url = "https://files.pythonhosted.org/packages/0b/78/1eeb63c3f9b2d1015a4c02788fb543141aad0a03ae3f7a7b669b2483f8d4/aiohttp-3.13.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7a653d872afe9f33497215745da7a943d1dc15b728a9c8da1c3ac423af35178e", size = 1792738, upload-time = "2025-10-28T20:58:29.787Z" },
- { url = "https://files.pythonhosted.org/packages/41/75/aaf1eea4c188e51538c04cc568040e3082db263a57086ea74a7d38c39e42/aiohttp-3.13.2-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:56d36e80d2003fa3fc0207fac644216d8532e9504a785ef9a8fd013f84a42c61", size = 1624061, upload-time = "2025-10-28T20:58:32.529Z" },
- { url = "https://files.pythonhosted.org/packages/9b/c2/3b6034de81fbcc43de8aeb209073a2286dfb50b86e927b4efd81cf848197/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:78cd586d8331fb8e241c2dd6b2f4061778cc69e150514b39a9e28dd050475661", size = 1789201, upload-time = "2025-10-28T20:58:34.618Z" },
- { url = "https://files.pythonhosted.org/packages/c9/38/c15dcf6d4d890217dae79d7213988f4e5fe6183d43893a9cf2fe9e84ca8d/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:20b10bbfbff766294fe99987f7bb3b74fdd2f1a2905f2562132641ad434dcf98", size = 1776868, upload-time = "2025-10-28T20:58:38.835Z" },
- { url = "https://files.pythonhosted.org/packages/04/75/f74fd178ac81adf4f283a74847807ade5150e48feda6aef024403716c30c/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:9ec49dff7e2b3c85cdeaa412e9d438f0ecd71676fde61ec57027dd392f00c693", size = 1790660, upload-time = "2025-10-28T20:58:41.507Z" },
- { url = "https://files.pythonhosted.org/packages/e7/80/7368bd0d06b16b3aba358c16b919e9c46cf11587dc572091031b0e9e3ef0/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:94f05348c4406450f9d73d38efb41d669ad6cd90c7ee194810d0eefbfa875a7a", size = 1617548, upload-time = "2025-10-28T20:58:43.674Z" },
- { url = "https://files.pythonhosted.org/packages/7d/4b/a6212790c50483cb3212e507378fbe26b5086d73941e1ec4b56a30439688/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:fa4dcb605c6f82a80c7f95713c2b11c3b8e9893b3ebd2bc9bde93165ed6107be", size = 1817240, upload-time = "2025-10-28T20:58:45.787Z" },
- { url = "https://files.pythonhosted.org/packages/ff/f7/ba5f0ba4ea8d8f3c32850912944532b933acbf0f3a75546b89269b9b7dde/aiohttp-3.13.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:cf00e5db968c3f67eccd2778574cf64d8b27d95b237770aa32400bd7a1ca4f6c", size = 1762334, upload-time = "2025-10-28T20:58:47.936Z" },
- { url = "https://files.pythonhosted.org/packages/7e/83/1a5a1856574588b1cad63609ea9ad75b32a8353ac995d830bf5da9357364/aiohttp-3.13.2-cp314-cp314t-win32.whl", hash = "sha256:d23b5fe492b0805a50d3371e8a728a9134d8de5447dce4c885f5587294750734", size = 464685, upload-time = "2025-10-28T20:58:50.642Z" },
- { url = "https://files.pythonhosted.org/packages/9f/4d/d22668674122c08f4d56972297c51a624e64b3ed1efaa40187607a7cb66e/aiohttp-3.13.2-cp314-cp314t-win_amd64.whl", hash = "sha256:ff0a7b0a82a7ab905cbda74006318d1b12e37c797eb1b0d4eb3e316cf47f658f", size = 498093, upload-time = "2025-10-28T20:58:52.782Z" },
+sdist = { url = "https://files.pythonhosted.org/packages/50/42/32cf8e7704ceb4481406eb87161349abb46a57fee3f008ba9cb610968646/aiohttp-3.13.3.tar.gz", hash = "sha256:a949eee43d3782f2daae4f4a2819b2cb9b0c5d3b7f7a927067cc84dafdbb9f88", size = 7844556, upload-time = "2026-01-03T17:33:05.204Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/36/d6/5aec9313ee6ea9c7cde8b891b69f4ff4001416867104580670a31daeba5b/aiohttp-3.13.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d5a372fd5afd301b3a89582817fdcdb6c34124787c70dbcc616f259013e7eef7", size = 738950, upload-time = "2026-01-03T17:29:13.002Z" },
+ { url = "https://files.pythonhosted.org/packages/68/03/8fa90a7e6d11ff20a18837a8e2b5dd23db01aabc475aa9271c8ad33299f5/aiohttp-3.13.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:147e422fd1223005c22b4fe080f5d93ced44460f5f9c105406b753612b587821", size = 496099, upload-time = "2026-01-03T17:29:15.268Z" },
+ { url = "https://files.pythonhosted.org/packages/d2/23/b81f744d402510a8366b74eb420fc0cc1170d0c43daca12d10814df85f10/aiohttp-3.13.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:859bd3f2156e81dd01432f5849fc73e2243d4a487c4fd26609b1299534ee1845", size = 491072, upload-time = "2026-01-03T17:29:16.922Z" },
+ { url = "https://files.pythonhosted.org/packages/d5/e1/56d1d1c0dd334cd203dd97706ce004c1aa24b34a813b0b8daf3383039706/aiohttp-3.13.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dca68018bf48c251ba17c72ed479f4dafe9dbd5a73707ad8d28a38d11f3d42af", size = 1671588, upload-time = "2026-01-03T17:29:18.539Z" },
+ { url = "https://files.pythonhosted.org/packages/5f/34/8d7f962604f4bc2b4e39eb1220dac7d4e4cba91fb9ba0474b4ecd67db165/aiohttp-3.13.3-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:fee0c6bc7db1de362252affec009707a17478a00ec69f797d23ca256e36d5940", size = 1640334, upload-time = "2026-01-03T17:29:21.028Z" },
+ { url = "https://files.pythonhosted.org/packages/94/1d/fcccf2c668d87337ddeef9881537baee13c58d8f01f12ba8a24215f2b804/aiohttp-3.13.3-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c048058117fd649334d81b4b526e94bde3ccaddb20463a815ced6ecbb7d11160", size = 1722656, upload-time = "2026-01-03T17:29:22.531Z" },
+ { url = "https://files.pythonhosted.org/packages/aa/98/c6f3b081c4c606bc1e5f2ec102e87d6411c73a9ef3616fea6f2d5c98c062/aiohttp-3.13.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:215a685b6fbbfcf71dfe96e3eba7a6f58f10da1dfdf4889c7dd856abe430dca7", size = 1817625, upload-time = "2026-01-03T17:29:24.276Z" },
+ { url = "https://files.pythonhosted.org/packages/2c/c0/cfcc3d2e11b477f86e1af2863f3858c8850d751ce8dc39c4058a072c9e54/aiohttp-3.13.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de2c184bb1fe2cbd2cefba613e9db29a5ab559323f994b6737e370d3da0ac455", size = 1672604, upload-time = "2026-01-03T17:29:26.099Z" },
+ { url = "https://files.pythonhosted.org/packages/1e/77/6b4ffcbcac4c6a5d041343a756f34a6dd26174ae07f977a64fe028dda5b0/aiohttp-3.13.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:75ca857eba4e20ce9f546cd59c7007b33906a4cd48f2ff6ccf1ccfc3b646f279", size = 1554370, upload-time = "2026-01-03T17:29:28.121Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/f0/e3ddfa93f17d689dbe014ba048f18e0c9f9b456033b70e94349a2e9048be/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:81e97251d9298386c2b7dbeb490d3d1badbdc69107fb8c9299dd04eb39bddc0e", size = 1642023, upload-time = "2026-01-03T17:29:30.002Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/45/c14019c9ec60a8e243d06d601b33dcc4fd92379424bde3021725859d7f99/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:c0e2d366af265797506f0283487223146af57815b388623f0357ef7eac9b209d", size = 1649680, upload-time = "2026-01-03T17:29:31.782Z" },
+ { url = "https://files.pythonhosted.org/packages/9c/fd/09c9451dae5aa5c5ed756df95ff9ef549d45d4be663bafd1e4954fd836f0/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4e239d501f73d6db1522599e14b9b321a7e3b1de66ce33d53a765d975e9f4808", size = 1692407, upload-time = "2026-01-03T17:29:33.392Z" },
+ { url = "https://files.pythonhosted.org/packages/a6/81/938bc2ec33c10efd6637ccb3d22f9f3160d08e8f3aa2587a2c2d5ab578eb/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:0db318f7a6f065d84cb1e02662c526294450b314a02bd9e2a8e67f0d8564ce40", size = 1543047, upload-time = "2026-01-03T17:29:34.855Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/23/80488ee21c8d567c83045e412e1d9b7077d27171591a4eb7822586e8c06a/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:bfc1cc2fe31a6026a8a88e4ecfb98d7f6b1fec150cfd708adbfd1d2f42257c29", size = 1715264, upload-time = "2026-01-03T17:29:36.389Z" },
+ { url = "https://files.pythonhosted.org/packages/e2/83/259a8da6683182768200b368120ab3deff5370bed93880fb9a3a86299f34/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:af71fff7bac6bb7508956696dce8f6eec2bbb045eceb40343944b1ae62b5ef11", size = 1657275, upload-time = "2026-01-03T17:29:38.162Z" },
+ { url = "https://files.pythonhosted.org/packages/3f/4f/2c41f800a0b560785c10fb316216ac058c105f9be50bdc6a285de88db625/aiohttp-3.13.3-cp310-cp310-win32.whl", hash = "sha256:37da61e244d1749798c151421602884db5270faf479cf0ef03af0ff68954c9dd", size = 434053, upload-time = "2026-01-03T17:29:40.074Z" },
+ { url = "https://files.pythonhosted.org/packages/80/df/29cd63c7ecfdb65ccc12f7d808cac4fa2a19544660c06c61a4a48462de0c/aiohttp-3.13.3-cp310-cp310-win_amd64.whl", hash = "sha256:7e63f210bc1b57ef699035f2b4b6d9ce096b5914414a49b0997c839b2bd2223c", size = 456687, upload-time = "2026-01-03T17:29:41.819Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/4c/a164164834f03924d9a29dc3acd9e7ee58f95857e0b467f6d04298594ebb/aiohttp-3.13.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:5b6073099fb654e0a068ae678b10feff95c5cae95bbfcbfa7af669d361a8aa6b", size = 746051, upload-time = "2026-01-03T17:29:43.287Z" },
+ { url = "https://files.pythonhosted.org/packages/82/71/d5c31390d18d4f58115037c432b7e0348c60f6f53b727cad33172144a112/aiohttp-3.13.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cb93e166e6c28716c8c6aeb5f99dfb6d5ccf482d29fe9bf9a794110e6d0ab64", size = 499234, upload-time = "2026-01-03T17:29:44.822Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/c9/741f8ac91e14b1d2e7100690425a5b2b919a87a5075406582991fb7de920/aiohttp-3.13.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:28e027cf2f6b641693a09f631759b4d9ce9165099d2b5d92af9bd4e197690eea", size = 494979, upload-time = "2026-01-03T17:29:46.405Z" },
+ { url = "https://files.pythonhosted.org/packages/75/b5/31d4d2e802dfd59f74ed47eba48869c1c21552c586d5e81a9d0d5c2ad640/aiohttp-3.13.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b61b7169ababd7802f9568ed96142616a9118dd2be0d1866e920e77ec8fa92a", size = 1748297, upload-time = "2026-01-03T17:29:48.083Z" },
+ { url = "https://files.pythonhosted.org/packages/1a/3e/eefad0ad42959f226bb79664826883f2687d602a9ae2941a18e0484a74d3/aiohttp-3.13.3-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:80dd4c21b0f6237676449c6baaa1039abae86b91636b6c91a7f8e61c87f89540", size = 1707172, upload-time = "2026-01-03T17:29:49.648Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/3a/54a64299fac2891c346cdcf2aa6803f994a2e4beeaf2e5a09dcc54acc842/aiohttp-3.13.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:65d2ccb7eabee90ce0503c17716fc77226be026dcc3e65cce859a30db715025b", size = 1805405, upload-time = "2026-01-03T17:29:51.244Z" },
+ { url = "https://files.pythonhosted.org/packages/6c/70/ddc1b7169cf64075e864f64595a14b147a895a868394a48f6a8031979038/aiohttp-3.13.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5b179331a481cb5529fca8b432d8d3c7001cb217513c94cd72d668d1248688a3", size = 1899449, upload-time = "2026-01-03T17:29:53.938Z" },
+ { url = "https://files.pythonhosted.org/packages/a1/7e/6815aab7d3a56610891c76ef79095677b8b5be6646aaf00f69b221765021/aiohttp-3.13.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d4c940f02f49483b18b079d1c27ab948721852b281f8b015c058100e9421dd1", size = 1748444, upload-time = "2026-01-03T17:29:55.484Z" },
+ { url = "https://files.pythonhosted.org/packages/6b/f2/073b145c4100da5511f457dc0f7558e99b2987cf72600d42b559db856fbc/aiohttp-3.13.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f9444f105664c4ce47a2a7171a2418bce5b7bae45fb610f4e2c36045d85911d3", size = 1606038, upload-time = "2026-01-03T17:29:57.179Z" },
+ { url = "https://files.pythonhosted.org/packages/0a/c1/778d011920cae03ae01424ec202c513dc69243cf2db303965615b81deeea/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:694976222c711d1d00ba131904beb60534f93966562f64440d0c9d41b8cdb440", size = 1724156, upload-time = "2026-01-03T17:29:58.914Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/cb/3419eabf4ec1e9ec6f242c32b689248365a1cf621891f6f0386632525494/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f33ed1a2bf1997a36661874b017f5c4b760f41266341af36febaf271d179f6d7", size = 1722340, upload-time = "2026-01-03T17:30:01.962Z" },
+ { url = "https://files.pythonhosted.org/packages/7a/e5/76cf77bdbc435bf233c1f114edad39ed4177ccbfab7c329482b179cff4f4/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:e636b3c5f61da31a92bf0d91da83e58fdfa96f178ba682f11d24f31944cdd28c", size = 1783041, upload-time = "2026-01-03T17:30:03.609Z" },
+ { url = "https://files.pythonhosted.org/packages/9d/d4/dd1ca234c794fd29c057ce8c0566b8ef7fd6a51069de5f06fa84b9a1971c/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:5d2d94f1f5fcbe40838ac51a6ab5704a6f9ea42e72ceda48de5e6b898521da51", size = 1596024, upload-time = "2026-01-03T17:30:05.132Z" },
+ { url = "https://files.pythonhosted.org/packages/55/58/4345b5f26661a6180afa686c473620c30a66afdf120ed3dd545bbc809e85/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:2be0e9ccf23e8a94f6f0650ce06042cefc6ac703d0d7ab6c7a917289f2539ad4", size = 1804590, upload-time = "2026-01-03T17:30:07.135Z" },
+ { url = "https://files.pythonhosted.org/packages/7b/06/05950619af6c2df7e0a431d889ba2813c9f0129cec76f663e547a5ad56f2/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9af5e68ee47d6534d36791bbe9b646d2a7c7deb6fc24d7943628edfbb3581f29", size = 1740355, upload-time = "2026-01-03T17:30:09.083Z" },
+ { url = "https://files.pythonhosted.org/packages/3e/80/958f16de79ba0422d7c1e284b2abd0c84bc03394fbe631d0a39ffa10e1eb/aiohttp-3.13.3-cp311-cp311-win32.whl", hash = "sha256:a2212ad43c0833a873d0fb3c63fa1bacedd4cf6af2fee62bf4b739ceec3ab239", size = 433701, upload-time = "2026-01-03T17:30:10.869Z" },
+ { url = "https://files.pythonhosted.org/packages/dc/f2/27cdf04c9851712d6c1b99df6821a6623c3c9e55956d4b1e318c337b5a48/aiohttp-3.13.3-cp311-cp311-win_amd64.whl", hash = "sha256:642f752c3eb117b105acbd87e2c143de710987e09860d674e068c4c2c441034f", size = 457678, upload-time = "2026-01-03T17:30:12.719Z" },
+ { url = "https://files.pythonhosted.org/packages/a0/be/4fc11f202955a69e0db803a12a062b8379c970c7c84f4882b6da17337cc1/aiohttp-3.13.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:b903a4dfee7d347e2d87697d0713be59e0b87925be030c9178c5faa58ea58d5c", size = 739732, upload-time = "2026-01-03T17:30:14.23Z" },
+ { url = "https://files.pythonhosted.org/packages/97/2c/621d5b851f94fa0bb7430d6089b3aa970a9d9b75196bc93bb624b0db237a/aiohttp-3.13.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a45530014d7a1e09f4a55f4f43097ba0fd155089372e105e4bff4ca76cb1b168", size = 494293, upload-time = "2026-01-03T17:30:15.96Z" },
+ { url = "https://files.pythonhosted.org/packages/5d/43/4be01406b78e1be8320bb8316dc9c42dbab553d281c40364e0f862d5661c/aiohttp-3.13.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:27234ef6d85c914f9efeb77ff616dbf4ad2380be0cda40b4db086ffc7ddd1b7d", size = 493533, upload-time = "2026-01-03T17:30:17.431Z" },
+ { url = "https://files.pythonhosted.org/packages/8d/a8/5a35dc56a06a2c90d4742cbf35294396907027f80eea696637945a106f25/aiohttp-3.13.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d32764c6c9aafb7fb55366a224756387cd50bfa720f32b88e0e6fa45b27dcf29", size = 1737839, upload-time = "2026-01-03T17:30:19.422Z" },
+ { url = "https://files.pythonhosted.org/packages/bf/62/4b9eeb331da56530bf2e198a297e5303e1c1ebdceeb00fe9b568a65c5a0c/aiohttp-3.13.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:b1a6102b4d3ebc07dad44fbf07b45bb600300f15b552ddf1851b5390202ea2e3", size = 1703932, upload-time = "2026-01-03T17:30:21.756Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/f6/af16887b5d419e6a367095994c0b1332d154f647e7dc2bd50e61876e8e3d/aiohttp-3.13.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c014c7ea7fb775dd015b2d3137378b7be0249a448a1612268b5a90c2d81de04d", size = 1771906, upload-time = "2026-01-03T17:30:23.932Z" },
+ { url = "https://files.pythonhosted.org/packages/ce/83/397c634b1bcc24292fa1e0c7822800f9f6569e32934bdeef09dae7992dfb/aiohttp-3.13.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2b8d8ddba8f95ba17582226f80e2de99c7a7948e66490ef8d947e272a93e9463", size = 1871020, upload-time = "2026-01-03T17:30:26Z" },
+ { url = "https://files.pythonhosted.org/packages/86/f6/a62cbbf13f0ac80a70f71b1672feba90fdb21fd7abd8dbf25c0105fb6fa3/aiohttp-3.13.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ae8dd55c8e6c4257eae3a20fd2c8f41edaea5992ed67156642493b8daf3cecc", size = 1755181, upload-time = "2026-01-03T17:30:27.554Z" },
+ { url = "https://files.pythonhosted.org/packages/0a/87/20a35ad487efdd3fba93d5843efdfaa62d2f1479eaafa7453398a44faf13/aiohttp-3.13.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01ad2529d4b5035578f5081606a465f3b814c542882804e2e8cda61adf5c71bf", size = 1561794, upload-time = "2026-01-03T17:30:29.254Z" },
+ { url = "https://files.pythonhosted.org/packages/de/95/8fd69a66682012f6716e1bc09ef8a1a2a91922c5725cb904689f112309c4/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bb4f7475e359992b580559e008c598091c45b5088f28614e855e42d39c2f1033", size = 1697900, upload-time = "2026-01-03T17:30:31.033Z" },
+ { url = "https://files.pythonhosted.org/packages/e5/66/7b94b3b5ba70e955ff597672dad1691333080e37f50280178967aff68657/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:c19b90316ad3b24c69cd78d5c9b4f3aa4497643685901185b65166293d36a00f", size = 1728239, upload-time = "2026-01-03T17:30:32.703Z" },
+ { url = "https://files.pythonhosted.org/packages/47/71/6f72f77f9f7d74719692ab65a2a0252584bf8d5f301e2ecb4c0da734530a/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:96d604498a7c782cb15a51c406acaea70d8c027ee6b90c569baa6e7b93073679", size = 1740527, upload-time = "2026-01-03T17:30:34.695Z" },
+ { url = "https://files.pythonhosted.org/packages/fa/b4/75ec16cbbd5c01bdaf4a05b19e103e78d7ce1ef7c80867eb0ace42ff4488/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:084911a532763e9d3dd95adf78a78f4096cd5f58cdc18e6fdbc1b58417a45423", size = 1554489, upload-time = "2026-01-03T17:30:36.864Z" },
+ { url = "https://files.pythonhosted.org/packages/52/8f/bc518c0eea29f8406dcf7ed1f96c9b48e3bc3995a96159b3fc11f9e08321/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:7a4a94eb787e606d0a09404b9c38c113d3b099d508021faa615d70a0131907ce", size = 1767852, upload-time = "2026-01-03T17:30:39.433Z" },
+ { url = "https://files.pythonhosted.org/packages/9d/f2/a07a75173124f31f11ea6f863dc44e6f09afe2bca45dd4e64979490deab1/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:87797e645d9d8e222e04160ee32aa06bc5c163e8499f24db719e7852ec23093a", size = 1722379, upload-time = "2026-01-03T17:30:41.081Z" },
+ { url = "https://files.pythonhosted.org/packages/3c/4a/1a3fee7c21350cac78e5c5cef711bac1b94feca07399f3d406972e2d8fcd/aiohttp-3.13.3-cp312-cp312-win32.whl", hash = "sha256:b04be762396457bef43f3597c991e192ee7da460a4953d7e647ee4b1c28e7046", size = 428253, upload-time = "2026-01-03T17:30:42.644Z" },
+ { url = "https://files.pythonhosted.org/packages/d9/b7/76175c7cb4eb73d91ad63c34e29fc4f77c9386bba4a65b53ba8e05ee3c39/aiohttp-3.13.3-cp312-cp312-win_amd64.whl", hash = "sha256:e3531d63d3bdfa7e3ac5e9b27b2dd7ec9df3206a98e0b3445fa906f233264c57", size = 455407, upload-time = "2026-01-03T17:30:44.195Z" },
+ { url = "https://files.pythonhosted.org/packages/97/8a/12ca489246ca1faaf5432844adbfce7ff2cc4997733e0af120869345643a/aiohttp-3.13.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:5dff64413671b0d3e7d5918ea490bdccb97a4ad29b3f311ed423200b2203e01c", size = 734190, upload-time = "2026-01-03T17:30:45.832Z" },
+ { url = "https://files.pythonhosted.org/packages/32/08/de43984c74ed1fca5c014808963cc83cb00d7bb06af228f132d33862ca76/aiohttp-3.13.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:87b9aab6d6ed88235aa2970294f496ff1a1f9adcd724d800e9b952395a80ffd9", size = 491783, upload-time = "2026-01-03T17:30:47.466Z" },
+ { url = "https://files.pythonhosted.org/packages/17/f8/8dd2cf6112a5a76f81f81a5130c57ca829d101ad583ce57f889179accdda/aiohttp-3.13.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:425c126c0dc43861e22cb1c14ba4c8e45d09516d0a3ae0a3f7494b79f5f233a3", size = 490704, upload-time = "2026-01-03T17:30:49.373Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/40/a46b03ca03936f832bc7eaa47cfbb1ad012ba1be4790122ee4f4f8cba074/aiohttp-3.13.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7f9120f7093c2a32d9647abcaf21e6ad275b4fbec5b55969f978b1a97c7c86bf", size = 1720652, upload-time = "2026-01-03T17:30:50.974Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/7e/917fe18e3607af92657e4285498f500dca797ff8c918bd7d90b05abf6c2a/aiohttp-3.13.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:697753042d57f4bf7122cab985bf15d0cef23c770864580f5af4f52023a56bd6", size = 1692014, upload-time = "2026-01-03T17:30:52.729Z" },
+ { url = "https://files.pythonhosted.org/packages/71/b6/cefa4cbc00d315d68973b671cf105b21a609c12b82d52e5d0c9ae61d2a09/aiohttp-3.13.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6de499a1a44e7de70735d0b39f67c8f25eb3d91eb3103be99ca0fa882cdd987d", size = 1759777, upload-time = "2026-01-03T17:30:54.537Z" },
+ { url = "https://files.pythonhosted.org/packages/fb/e3/e06ee07b45e59e6d81498b591fc589629be1553abb2a82ce33efe2a7b068/aiohttp-3.13.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:37239e9f9a7ea9ac5bf6b92b0260b01f8a22281996da609206a84df860bc1261", size = 1861276, upload-time = "2026-01-03T17:30:56.512Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/24/75d274228acf35ceeb2850b8ce04de9dd7355ff7a0b49d607ee60c29c518/aiohttp-3.13.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f76c1e3fe7d7c8afad7ed193f89a292e1999608170dcc9751a7462a87dfd5bc0", size = 1743131, upload-time = "2026-01-03T17:30:58.256Z" },
+ { url = "https://files.pythonhosted.org/packages/04/98/3d21dde21889b17ca2eea54fdcff21b27b93f45b7bb94ca029c31ab59dc3/aiohttp-3.13.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fc290605db2a917f6e81b0e1e0796469871f5af381ce15c604a3c5c7e51cb730", size = 1556863, upload-time = "2026-01-03T17:31:00.445Z" },
+ { url = "https://files.pythonhosted.org/packages/9e/84/da0c3ab1192eaf64782b03971ab4055b475d0db07b17eff925e8c93b3aa5/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4021b51936308aeea0367b8f006dc999ca02bc118a0cc78c303f50a2ff6afb91", size = 1682793, upload-time = "2026-01-03T17:31:03.024Z" },
+ { url = "https://files.pythonhosted.org/packages/ff/0f/5802ada182f575afa02cbd0ec5180d7e13a402afb7c2c03a9aa5e5d49060/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:49a03727c1bba9a97d3e93c9f93ca03a57300f484b6e935463099841261195d3", size = 1716676, upload-time = "2026-01-03T17:31:04.842Z" },
+ { url = "https://files.pythonhosted.org/packages/3f/8c/714d53bd8b5a4560667f7bbbb06b20c2382f9c7847d198370ec6526af39c/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3d9908a48eb7416dc1f4524e69f1d32e5d90e3981e4e37eb0aa1cd18f9cfa2a4", size = 1733217, upload-time = "2026-01-03T17:31:06.868Z" },
+ { url = "https://files.pythonhosted.org/packages/7d/79/e2176f46d2e963facea939f5be2d26368ce543622be6f00a12844d3c991f/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2712039939ec963c237286113c68dbad80a82a4281543f3abf766d9d73228998", size = 1552303, upload-time = "2026-01-03T17:31:08.958Z" },
+ { url = "https://files.pythonhosted.org/packages/ab/6a/28ed4dea1759916090587d1fe57087b03e6c784a642b85ef48217b0277ae/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:7bfdc049127717581866fa4708791220970ce291c23e28ccf3922c700740fdc0", size = 1763673, upload-time = "2026-01-03T17:31:10.676Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/35/4a3daeb8b9fab49240d21c04d50732313295e4bd813a465d840236dd0ce1/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8057c98e0c8472d8846b9c79f56766bcc57e3e8ac7bfd510482332366c56c591", size = 1721120, upload-time = "2026-01-03T17:31:12.575Z" },
+ { url = "https://files.pythonhosted.org/packages/bc/9f/d643bb3c5fb99547323e635e251c609fbbc660d983144cfebec529e09264/aiohttp-3.13.3-cp313-cp313-win32.whl", hash = "sha256:1449ceddcdbcf2e0446957863af03ebaaa03f94c090f945411b61269e2cb5daf", size = 427383, upload-time = "2026-01-03T17:31:14.382Z" },
+ { url = "https://files.pythonhosted.org/packages/4e/f1/ab0395f8a79933577cdd996dd2f9aa6014af9535f65dddcf88204682fe62/aiohttp-3.13.3-cp313-cp313-win_amd64.whl", hash = "sha256:693781c45a4033d31d4187d2436f5ac701e7bbfe5df40d917736108c1cc7436e", size = 453899, upload-time = "2026-01-03T17:31:15.958Z" },
+ { url = "https://files.pythonhosted.org/packages/99/36/5b6514a9f5d66f4e2597e40dea2e3db271e023eb7a5d22defe96ba560996/aiohttp-3.13.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:ea37047c6b367fd4bd632bff8077449b8fa034b69e812a18e0132a00fae6e808", size = 737238, upload-time = "2026-01-03T17:31:17.909Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/49/459327f0d5bcd8c6c9ca69e60fdeebc3622861e696490d8674a6d0cb90a6/aiohttp-3.13.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:6fc0e2337d1a4c3e6acafda6a78a39d4c14caea625124817420abceed36e2415", size = 492292, upload-time = "2026-01-03T17:31:19.919Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/0b/b97660c5fd05d3495b4eb27f2d0ef18dc1dc4eff7511a9bf371397ff0264/aiohttp-3.13.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c685f2d80bb67ca8c3837823ad76196b3694b0159d232206d1e461d3d434666f", size = 493021, upload-time = "2026-01-03T17:31:21.636Z" },
+ { url = "https://files.pythonhosted.org/packages/54/d4/438efabdf74e30aeceb890c3290bbaa449780583b1270b00661126b8aae4/aiohttp-3.13.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48e377758516d262bde50c2584fc6c578af272559c409eecbdd2bae1601184d6", size = 1717263, upload-time = "2026-01-03T17:31:23.296Z" },
+ { url = "https://files.pythonhosted.org/packages/71/f2/7bddc7fd612367d1459c5bcf598a9e8f7092d6580d98de0e057eb42697ad/aiohttp-3.13.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:34749271508078b261c4abb1767d42b8d0c0cc9449c73a4df494777dc55f0687", size = 1669107, upload-time = "2026-01-03T17:31:25.334Z" },
+ { url = "https://files.pythonhosted.org/packages/00/5a/1aeaecca40e22560f97610a329e0e5efef5e0b5afdf9f857f0d93839ab2e/aiohttp-3.13.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:82611aeec80eb144416956ec85b6ca45a64d76429c1ed46ae1b5f86c6e0c9a26", size = 1760196, upload-time = "2026-01-03T17:31:27.394Z" },
+ { url = "https://files.pythonhosted.org/packages/f8/f8/0ff6992bea7bd560fc510ea1c815f87eedd745fe035589c71ce05612a19a/aiohttp-3.13.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2fff83cfc93f18f215896e3a190e8e5cb413ce01553901aca925176e7568963a", size = 1843591, upload-time = "2026-01-03T17:31:29.238Z" },
+ { url = "https://files.pythonhosted.org/packages/e3/d1/e30e537a15f53485b61f5be525f2157da719819e8377298502aebac45536/aiohttp-3.13.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bbe7d4cecacb439e2e2a8a1a7b935c25b812af7a5fd26503a66dadf428e79ec1", size = 1720277, upload-time = "2026-01-03T17:31:31.053Z" },
+ { url = "https://files.pythonhosted.org/packages/84/45/23f4c451d8192f553d38d838831ebbc156907ea6e05557f39563101b7717/aiohttp-3.13.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b928f30fe49574253644b1ca44b1b8adbd903aa0da4b9054a6c20fc7f4092a25", size = 1548575, upload-time = "2026-01-03T17:31:32.87Z" },
+ { url = "https://files.pythonhosted.org/packages/6a/ed/0a42b127a43712eda7807e7892c083eadfaf8429ca8fb619662a530a3aab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7b5e8fe4de30df199155baaf64f2fcd604f4c678ed20910db8e2c66dc4b11603", size = 1679455, upload-time = "2026-01-03T17:31:34.76Z" },
+ { url = "https://files.pythonhosted.org/packages/2e/b5/c05f0c2b4b4fe2c9d55e73b6d3ed4fd6c9dc2684b1d81cbdf77e7fad9adb/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:8542f41a62bcc58fc7f11cf7c90e0ec324ce44950003feb70640fc2a9092c32a", size = 1687417, upload-time = "2026-01-03T17:31:36.699Z" },
+ { url = "https://files.pythonhosted.org/packages/c9/6b/915bc5dad66aef602b9e459b5a973529304d4e89ca86999d9d75d80cbd0b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:5e1d8c8b8f1d91cd08d8f4a3c2b067bfca6ec043d3ff36de0f3a715feeedf926", size = 1729968, upload-time = "2026-01-03T17:31:38.622Z" },
+ { url = "https://files.pythonhosted.org/packages/11/3b/e84581290a9520024a08640b63d07673057aec5ca548177a82026187ba73/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:90455115e5da1c3c51ab619ac57f877da8fd6d73c05aacd125c5ae9819582aba", size = 1545690, upload-time = "2026-01-03T17:31:40.57Z" },
+ { url = "https://files.pythonhosted.org/packages/f5/04/0c3655a566c43fd647c81b895dfe361b9f9ad6d58c19309d45cff52d6c3b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:042e9e0bcb5fba81886c8b4fbb9a09d6b8a00245fd8d88e4d989c1f96c74164c", size = 1746390, upload-time = "2026-01-03T17:31:42.857Z" },
+ { url = "https://files.pythonhosted.org/packages/1f/53/71165b26978f719c3419381514c9690bd5980e764a09440a10bb816ea4ab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2eb752b102b12a76ca02dff751a801f028b4ffbbc478840b473597fc91a9ed43", size = 1702188, upload-time = "2026-01-03T17:31:44.984Z" },
+ { url = "https://files.pythonhosted.org/packages/29/a7/cbe6c9e8e136314fa1980da388a59d2f35f35395948a08b6747baebb6aa6/aiohttp-3.13.3-cp314-cp314-win32.whl", hash = "sha256:b556c85915d8efaed322bf1bdae9486aa0f3f764195a0fb6ee962e5c71ef5ce1", size = 433126, upload-time = "2026-01-03T17:31:47.463Z" },
+ { url = "https://files.pythonhosted.org/packages/de/56/982704adea7d3b16614fc5936014e9af85c0e34b58f9046655817f04306e/aiohttp-3.13.3-cp314-cp314-win_amd64.whl", hash = "sha256:9bf9f7a65e7aa20dd764151fb3d616c81088f91f8df39c3893a536e279b4b984", size = 459128, upload-time = "2026-01-03T17:31:49.2Z" },
+ { url = "https://files.pythonhosted.org/packages/6c/2a/3c79b638a9c3d4658d345339d22070241ea341ed4e07b5ac60fb0f418003/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:05861afbbec40650d8a07ea324367cb93e9e8cc7762e04dd4405df99fa65159c", size = 769512, upload-time = "2026-01-03T17:31:51.134Z" },
+ { url = "https://files.pythonhosted.org/packages/29/b9/3e5014d46c0ab0db8707e0ac2711ed28c4da0218c358a4e7c17bae0d8722/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2fc82186fadc4a8316768d61f3722c230e2c1dcab4200d52d2ebdf2482e47592", size = 506444, upload-time = "2026-01-03T17:31:52.85Z" },
+ { url = "https://files.pythonhosted.org/packages/90/03/c1d4ef9a054e151cd7839cdc497f2638f00b93cbe8043983986630d7a80c/aiohttp-3.13.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0add0900ff220d1d5c5ebbf99ed88b0c1bbf87aa7e4262300ed1376a6b13414f", size = 510798, upload-time = "2026-01-03T17:31:54.91Z" },
+ { url = "https://files.pythonhosted.org/packages/ea/76/8c1e5abbfe8e127c893fe7ead569148a4d5a799f7cf958d8c09f3eedf097/aiohttp-3.13.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:568f416a4072fbfae453dcf9a99194bbb8bdeab718e08ee13dfa2ba0e4bebf29", size = 1868835, upload-time = "2026-01-03T17:31:56.733Z" },
+ { url = "https://files.pythonhosted.org/packages/8e/ac/984c5a6f74c363b01ff97adc96a3976d9c98940b8969a1881575b279ac5d/aiohttp-3.13.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:add1da70de90a2569c5e15249ff76a631ccacfe198375eead4aadf3b8dc849dc", size = 1720486, upload-time = "2026-01-03T17:31:58.65Z" },
+ { url = "https://files.pythonhosted.org/packages/b2/9a/b7039c5f099c4eb632138728828b33428585031a1e658d693d41d07d89d1/aiohttp-3.13.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:10b47b7ba335d2e9b1239fa571131a87e2d8ec96b333e68b2a305e7a98b0bae2", size = 1847951, upload-time = "2026-01-03T17:32:00.989Z" },
+ { url = "https://files.pythonhosted.org/packages/3c/02/3bec2b9a1ba3c19ff89a43a19324202b8eb187ca1e928d8bdac9bbdddebd/aiohttp-3.13.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4dce1c718e38081c8f35f323209d4c1df7d4db4bab1b5c88a6b4d12b74587", size = 1941001, upload-time = "2026-01-03T17:32:03.122Z" },
+ { url = "https://files.pythonhosted.org/packages/37/df/d879401cedeef27ac4717f6426c8c36c3091c6e9f08a9178cc87549c537f/aiohttp-3.13.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34bac00a67a812570d4a460447e1e9e06fae622946955f939051e7cc895cfab8", size = 1797246, upload-time = "2026-01-03T17:32:05.255Z" },
+ { url = "https://files.pythonhosted.org/packages/8d/15/be122de1f67e6953add23335c8ece6d314ab67c8bebb3f181063010795a7/aiohttp-3.13.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a19884d2ee70b06d9204b2727a7b9f983d0c684c650254679e716b0b77920632", size = 1627131, upload-time = "2026-01-03T17:32:07.607Z" },
+ { url = "https://files.pythonhosted.org/packages/12/12/70eedcac9134cfa3219ab7af31ea56bc877395b1ac30d65b1bc4b27d0438/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5f8ca7f2bb6ba8348a3614c7918cc4bb73268c5ac2a207576b7afea19d3d9f64", size = 1795196, upload-time = "2026-01-03T17:32:09.59Z" },
+ { url = "https://files.pythonhosted.org/packages/32/11/b30e1b1cd1f3054af86ebe60df96989c6a414dd87e27ad16950eee420bea/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:b0d95340658b9d2f11d9697f59b3814a9d3bb4b7a7c20b131df4bcef464037c0", size = 1782841, upload-time = "2026-01-03T17:32:11.445Z" },
+ { url = "https://files.pythonhosted.org/packages/88/0d/d98a9367b38912384a17e287850f5695c528cff0f14f791ce8ee2e4f7796/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:a1e53262fd202e4b40b70c3aff944a8155059beedc8a89bba9dc1f9ef06a1b56", size = 1795193, upload-time = "2026-01-03T17:32:13.705Z" },
+ { url = "https://files.pythonhosted.org/packages/43/a5/a2dfd1f5ff5581632c7f6a30e1744deda03808974f94f6534241ef60c751/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:d60ac9663f44168038586cab2157e122e46bdef09e9368b37f2d82d354c23f72", size = 1621979, upload-time = "2026-01-03T17:32:15.965Z" },
+ { url = "https://files.pythonhosted.org/packages/fa/f0/12973c382ae7c1cccbc4417e129c5bf54c374dfb85af70893646e1f0e749/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:90751b8eed69435bac9ff4e3d2f6b3af1f57e37ecb0fbeee59c0174c9e2d41df", size = 1822193, upload-time = "2026-01-03T17:32:18.219Z" },
+ { url = "https://files.pythonhosted.org/packages/3c/5f/24155e30ba7f8c96918af1350eb0663e2430aad9e001c0489d89cd708ab1/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fc353029f176fd2b3ec6cfc71be166aba1936fe5d73dd1992ce289ca6647a9aa", size = 1769801, upload-time = "2026-01-03T17:32:20.25Z" },
+ { url = "https://files.pythonhosted.org/packages/eb/f8/7314031ff5c10e6ece114da79b338ec17eeff3a079e53151f7e9f43c4723/aiohttp-3.13.3-cp314-cp314t-win32.whl", hash = "sha256:2e41b18a58da1e474a057b3d35248d8320029f61d70a37629535b16a0c8f3767", size = 466523, upload-time = "2026-01-03T17:32:22.215Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/63/278a98c715ae467624eafe375542d8ba9b4383a016df8fdefe0ae28382a7/aiohttp-3.13.3-cp314-cp314t-win_amd64.whl", hash = "sha256:44531a36aa2264a1860089ffd4dce7baf875ee5a6079d5fb42e261c704ef7344", size = 499694, upload-time = "2026-01-03T17:32:24.546Z" },
]
[[package]]
@@ -3455,11 +3455,11 @@ wheels = [
[[package]]
name = "urllib3"
-version = "2.6.2"
+version = "2.6.3"
source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1e/24/a2a2ed9addd907787d7aa0355ba36a6cadf1768b934c652ea78acbd59dcd/urllib3-2.6.2.tar.gz", hash = "sha256:016f9c98bb7e98085cb2b4b17b87d2c702975664e4f060c6532e64d1c1a5e797", size = 432930, upload-time = "2025-12-11T15:56:40.252Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/6d/b9/4095b668ea3678bf6a0af005527f39de12fb026516fb3df17495a733b7f8/urllib3-2.6.2-py3-none-any.whl", hash = "sha256:ec21cddfe7724fc7cb4ba4bea7aa8e2ef36f607a4bab81aa6ce42a13dc3f03dd", size = 131182, upload-time = "2025-12-11T15:56:38.584Z" },
+ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" },
]
[[package]]
From d732e13ac0f50367b84b6d40d111b6b9e2334c9d Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 09:21:21 -0800
Subject: [PATCH 50/60] fix typo
---
README.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/README.md b/README.md
index c5f1f10b..ca0a8dfc 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,7 @@ Requires Python 3.10 or newer. For development setup, see [CONTRIBUTING.md](CONT
- **Python 3.10+**
- **An LLM API key** (e.g., OpenAI, Anthropic, Groq) -- required for `osmosis test` and `osmosis eval`. See [supported providers](https://docs.litellm.ai/docs/providers).
-- **Osmosis account** (optional) -- needed for platform features like `osmosis login`, workspace management, and submitting training runs. Sign up at [osmosis.ai](https://osmosis.ai).
+- **Osmosis account** (optional) -- needed for platform features like `osmosis login`, workspace management, and submitting training runs. Sign up at [platform.osmosis.ai](https://platform.osmosis.ai).
### pip
From e968849e89fffab48ff56c2e380fcf0b60cd5532 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 09:33:20 -0800
Subject: [PATCH 51/60] Enhance documentation clarity and structure by refining
sections and improving examples. Updated README to better outline Local and
Remote Rollout modes, ensuring consistency in explanations and error
handling.
---
.github/workflows/auto-label-pr.yml | 63 +++++++++++++++++++++++++++++
1 file changed, 63 insertions(+)
create mode 100644 .github/workflows/auto-label-pr.yml
diff --git a/.github/workflows/auto-label-pr.yml b/.github/workflows/auto-label-pr.yml
new file mode 100644
index 00000000..1d65b434
--- /dev/null
+++ b/.github/workflows/auto-label-pr.yml
@@ -0,0 +1,63 @@
+name: Auto Label PR
+
+on:
+ pull_request:
+ types: [opened, edited]
+
+jobs:
+ label:
+ runs-on: ubuntu-latest
+ permissions:
+ pull-requests: write
+ steps:
+ - name: Apply labels based on PR title
+ env:
+ GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+ PR_NUMBER: ${{ github.event.pull_request.number }}
+ PR_TITLE: ${{ github.event.pull_request.title }}
+ run: |
+ labels=()
+
+ # Breaking change
+ if [[ "$PR_TITLE" =~ ^\[BREAKING\] ]]; then
+ labels+=("breaking")
+ fi
+
+ # Type-based labels (maps to release.yml categories)
+ if [[ "$PR_TITLE" =~ \]\ feat: ]]; then
+ labels+=("enhancement")
+ elif [[ "$PR_TITLE" =~ \]\ fix: ]]; then
+ labels+=("bug")
+ elif [[ "$PR_TITLE" =~ \]\ doc: ]]; then
+ labels+=("documentation")
+ elif [[ "$PR_TITLE" =~ \]\ chore: ]]; then
+ labels+=("chore")
+ elif [[ "$PR_TITLE" =~ \]\ refactor: ]]; then
+ labels+=("refactor")
+ elif [[ "$PR_TITLE" =~ \]\ test: ]]; then
+ labels+=("chore")
+ fi
+
+ # Module label from [module] in title
+ if [[ "$PR_TITLE" =~ \[([a-z]+)\] ]]; then
+ module="${BASH_REMATCH[1]}"
+ # Map [doc] module to existing 'documentation' label
+ if [[ "$module" == "doc" ]]; then
+ labels+=("documentation")
+ elif [[ "$module" != "misc" ]]; then
+ labels+=("$module")
+ fi
+ fi
+
+ if [ ${#labels[@]} -eq 0 ]; then
+ echo "No labels to apply"
+ exit 0
+ fi
+
+ # Deduplicate
+ labels=($(printf '%s\n' "${labels[@]}" | sort -u))
+
+ echo "Applying labels: ${labels[*]}"
+ gh pr edit "$PR_NUMBER" \
+ --repo "${{ github.repository }}" \
+ --add-label "$(IFS=,; echo "${labels[*]}")"
From 68fdc6913b740a0f8993ef6d5121b874d7095dde Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 09:46:59 -0800
Subject: [PATCH 52/60] cleanup some dead code
---
osmosis_ai/auth/credentials.py | 4 ----
osmosis_ai/rollout/validator.py | 5 +----
2 files changed, 1 insertion(+), 8 deletions(-)
diff --git a/osmosis_ai/auth/credentials.py b/osmosis_ai/auth/credentials.py
index fe127c4b..529b4b39 100644
--- a/osmosis_ai/auth/credentials.py
+++ b/osmosis_ai/auth/credentials.py
@@ -136,10 +136,6 @@ def get_active_credentials(self) -> WorkspaceCredentials | None:
return None
return self.workspaces.get(self.active_workspace)
- def get_workspace_names(self) -> list[str]:
- """Get list of all workspace names."""
- return list(self.workspaces.keys())
-
def _ensure_config_dir() -> None:
"""Ensure the config directory exists with proper permissions."""
diff --git a/osmosis_ai/rollout/validator.py b/osmosis_ai/rollout/validator.py
index 943390a9..4e587ba2 100644
--- a/osmosis_ai/rollout/validator.py
+++ b/osmosis_ai/rollout/validator.py
@@ -24,7 +24,7 @@ class MyAgent(RolloutAgentLoop):
import logging
from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
+from typing import Any
from osmosis_ai.rollout.core.base import RolloutAgentLoop
from osmosis_ai.rollout.core.schemas import (
@@ -32,9 +32,6 @@ class MyAgent(RolloutAgentLoop):
RolloutRequest,
)
-if TYPE_CHECKING:
- pass
-
logger = logging.getLogger(__name__)
From f30b886d21e0571256332403c2f7eaa406d04114 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 10:05:23 -0800
Subject: [PATCH 53/60] cleanup
---
examples/reward_functions.py | 25 --------------------
examples/rubric_functions.py | 1 -
osmosis_ai/rollout/eval/evaluation/cli.py | 3 ---
osmosis_ai/rollout/eval/evaluation/runner.py | 2 +-
osmosis_ai/rollout/eval/test_mode/cli.py | 3 ---
osmosis_ai/rollout/server/api_key.py | 3 ---
osmosis_ai/rollout/server/serve.py | 1 -
osmosis_ai/rollout/tools.py | 3 ---
osmosis_ai/rollout/validator.py | 3 ---
9 files changed, 1 insertion(+), 43 deletions(-)
diff --git a/examples/reward_functions.py b/examples/reward_functions.py
index 6a269148..4e586f27 100644
--- a/examples/reward_functions.py
+++ b/examples/reward_functions.py
@@ -66,31 +66,6 @@ def minimal_reward(
return float(solution_str == ground_truth)
-# INCORRECT USAGE EXAMPLES (these will raise TypeError)
-
-# Uncomment these to see the validation errors:
-
-# @osmosis_reward # Wrong parameter names
-# def wrong_names(input_str: str, expected: str, extra: dict = None):
-# return input_str == expected
-
-# @osmosis_reward # Missing type annotations
-# def missing_annotations(solution_str, ground_truth, extra_info=None):
-# return solution_str == ground_truth
-
-# @osmosis_reward # Wrong type annotations
-# def wrong_types(solution_str: int, ground_truth: str, extra_info: dict | None = None):
-# return str(solution_str) == ground_truth
-
-# @osmosis_reward # Too many parameters
-# def too_many_params(solution_str: str, ground_truth: str, extra_info: dict | None = None, bonus: float = 0.0):
-# return (solution_str == ground_truth) + bonus
-
-# @osmosis_reward # Missing default value for extra_info
-# def no_default(solution_str: str, ground_truth: str, extra_info: dict):
-# return solution_str == ground_truth
-
-
if __name__ == "__main__":
# Test the reward functions
test_cases = [
diff --git a/examples/rubric_functions.py b/examples/rubric_functions.py
index 3d7a2346..095f5962 100644
--- a/examples/rubric_functions.py
+++ b/examples/rubric_functions.py
@@ -253,4 +253,3 @@ def run_cerebras_example() -> None:
# run_xai_example()
run_openrouter_example()
run_cerebras_example()
- pass
diff --git a/osmosis_ai/rollout/eval/evaluation/cli.py b/osmosis_ai/rollout/eval/evaluation/cli.py
index 30a485bb..c84f61bb 100644
--- a/osmosis_ai/rollout/eval/evaluation/cli.py
+++ b/osmosis_ai/rollout/eval/evaluation/cli.py
@@ -28,9 +28,6 @@
from osmosis_ai.rollout.eval.evaluation.runner import EvalResult, EvalRunResult
-logger = logging.getLogger(__name__)
-
-
class EvalCommand:
"""Handler for `osmosis eval`."""
diff --git a/osmosis_ai/rollout/eval/evaluation/runner.py b/osmosis_ai/rollout/eval/evaluation/runner.py
index 4c665384..f5e5b67f 100644
--- a/osmosis_ai/rollout/eval/evaluation/runner.py
+++ b/osmosis_ai/rollout/eval/evaluation/runner.py
@@ -621,7 +621,7 @@ async def _run_one(
try:
for task in asyncio.as_completed(tasks):
try:
- _row_index, _result = await task
+ await task
except SystemicProviderError as e:
# Record error but keep awaiting remaining
# tasks in this batch so their results are collected.
diff --git a/osmosis_ai/rollout/eval/test_mode/cli.py b/osmosis_ai/rollout/eval/test_mode/cli.py
index 125751a2..cedc5d8d 100644
--- a/osmosis_ai/rollout/eval/test_mode/cli.py
+++ b/osmosis_ai/rollout/eval/test_mode/cli.py
@@ -35,9 +35,6 @@
)
-logger = logging.getLogger(__name__)
-
-
@dataclass
class _SetupResult:
"""Result of setup phase containing initialized components."""
diff --git a/osmosis_ai/rollout/server/api_key.py b/osmosis_ai/rollout/server/api_key.py
index 2755865f..ab57930a 100644
--- a/osmosis_ai/rollout/server/api_key.py
+++ b/osmosis_ai/rollout/server/api_key.py
@@ -6,11 +6,8 @@
from __future__ import annotations
-import logging
import secrets
-logger = logging.getLogger(__name__)
-
# API Key prefix for identification
API_KEY_PREFIX = "osm_rollout_"
diff --git a/osmosis_ai/rollout/server/serve.py b/osmosis_ai/rollout/server/serve.py
index 8d9ebe66..2af22a4a 100644
--- a/osmosis_ai/rollout/server/serve.py
+++ b/osmosis_ai/rollout/server/serve.py
@@ -159,7 +159,6 @@ def serve_agent_loop(
if local_debug:
logger.info("Local debug mode enabled: API key authentication disabled")
api_key = None
- api_key_provided = False
else:
# Generate API Key for this server instance if not provided
if api_key is None:
diff --git a/osmosis_ai/rollout/tools.py b/osmosis_ai/rollout/tools.py
index c44b5348..cb78f5c3 100644
--- a/osmosis_ai/rollout/tools.py
+++ b/osmosis_ai/rollout/tools.py
@@ -31,14 +31,11 @@ async def my_executor(tool_call):
import asyncio
import json
-import logging
from collections.abc import Awaitable, Callable
from typing import Any
from osmosis_ai.rollout.core.exceptions import ToolArgumentError
-logger = logging.getLogger(__name__)
-
def create_tool_result(tool_call_id: str, content: str) -> dict[str, str]:
"""Create a standardized tool result message.
diff --git a/osmosis_ai/rollout/validator.py b/osmosis_ai/rollout/validator.py
index 4e587ba2..63825c07 100644
--- a/osmosis_ai/rollout/validator.py
+++ b/osmosis_ai/rollout/validator.py
@@ -22,7 +22,6 @@ class MyAgent(RolloutAgentLoop):
from __future__ import annotations
-import logging
from dataclasses import dataclass, field
from typing import Any
@@ -32,8 +31,6 @@ class MyAgent(RolloutAgentLoop):
RolloutRequest,
)
-logger = logging.getLogger(__name__)
-
@dataclass
class ValidationError:
From e4f7ff542de67a45f0ff326f5986a4edff33c84f Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 10:27:09 -0800
Subject: [PATCH 54/60] Enhance test coverage reporting in GitHub Actions by
adding XML output and integrating Codecov for coverage uploads on Python
3.12.
---
.github/workflows/tests.yml | 9 ++++++++-
codecov.yml | 13 +++++++++++++
2 files changed, 21 insertions(+), 1 deletion(-)
create mode 100644 codecov.yml
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index b5b30c56..11a06739 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -57,7 +57,14 @@ jobs:
- name: Run tests with coverage
if: matrix.python-version == '3.12'
- run: uv run pytest --cov=osmosis_ai --cov-report=term-missing
+ run: uv run pytest --cov=osmosis_ai --cov-branch --cov-report=term-missing --cov-report=xml
+
+ - name: Upload coverage to Codecov
+ if: matrix.python-version == '3.12'
+ uses: codecov/codecov-action@v5
+ with:
+ files: coverage.xml
+ token: ${{ secrets.CODECOV_TOKEN }}
- name: Run tests
if: matrix.python-version != '3.12'
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 00000000..c8986c58
--- /dev/null
+++ b/codecov.yml
@@ -0,0 +1,13 @@
+coverage:
+ status:
+ project:
+ default:
+ target: 70%
+ threshold: 2%
+ patch:
+ default:
+ target: 70%
+
+comment:
+ layout: "diff, flags, files"
+ require_changes: true
From 8c4847ce7503636aa2044814969c34473af89abe Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 10:35:44 -0800
Subject: [PATCH 55/60] Update GitHub Actions workflow to conditionally upload
coverage to Codecov only for non-fork pull requests on Python 3.12.
---
.github/workflows/tests.yml | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 11a06739..545c88e3 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -60,7 +60,7 @@ jobs:
run: uv run pytest --cov=osmosis_ai --cov-branch --cov-report=term-missing --cov-report=xml
- name: Upload coverage to Codecov
- if: matrix.python-version == '3.12'
+ if: matrix.python-version == '3.12' && !github.event.pull_request.head.repo.fork
uses: codecov/codecov-action@v5
with:
files: coverage.xml
From d401fa52a490439505074bd52759f4991b2c7125 Mon Sep 17 00:00:00 2001
From: Brian <122925040+JoyboyBrian@users.noreply.github.com>
Date: Wed, 18 Feb 2026 11:18:22 -0800
Subject: [PATCH 56/60] Add Codecov badge to README
Added Codecov badge to README for coverage tracking.
---
README.md | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/README.md b/README.md
index ca0a8dfc..0b97bb88 100644
--- a/README.md
+++ b/README.md
@@ -10,10 +10,14 @@
+
+
+
+
# osmosis-ai
> ⚠️ **Warning**: osmosis-ai is still in active development. APIs may change between versions.
From 557c22471d23e9d76b951ed9c4435eb8c724c7b6 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 11:41:21 -0800
Subject: [PATCH 57/60] Refactor warning handling in CLI and remove unnecessary
filters from rubric_eval and llm_client modules. Suppress Python warnings in
CLI for a cleaner user experience.
---
osmosis_ai/cli.py | 6 ++++++
osmosis_ai/rollout/eval/common/llm_client.py | 7 -------
osmosis_ai/rubric_eval.py | 13 -------------
3 files changed, 6 insertions(+), 20 deletions(-)
diff --git a/osmosis_ai/cli.py b/osmosis_ai/cli.py
index 4c922183..85da8149 100644
--- a/osmosis_ai/cli.py
+++ b/osmosis_ai/cli.py
@@ -2,6 +2,7 @@
import argparse
import sys
+import warnings
from dotenv import load_dotenv
@@ -19,6 +20,11 @@
def main(argv: list[str] | None = None) -> int:
"""Entry point for the osmosis CLI."""
+ # Suppress all Python warnings for a clean CLI experience.
+ # This only affects the CLI process; library consumers are not impacted.
+ # Meaningful user-facing messages should use logging or print instead.
+ warnings.filterwarnings("ignore")
+
# Load environment variables from .env file in current working directory
load_dotenv()
diff --git a/osmosis_ai/rollout/eval/common/llm_client.py b/osmosis_ai/rollout/eval/common/llm_client.py
index fe7abfcc..b63309f7 100644
--- a/osmosis_ai/rollout/eval/common/llm_client.py
+++ b/osmosis_ai/rollout/eval/common/llm_client.py
@@ -8,15 +8,8 @@
import inspect
import logging
import time
-import warnings
from typing import Any
-warnings.filterwarnings(
- "ignore",
- message="Pydantic serializer warnings:",
- category=UserWarning,
-)
-
from osmosis_ai._litellm_compat import APIConnectionError as _APIConnectionError
from osmosis_ai._litellm_compat import AuthenticationError as _AuthenticationError
from osmosis_ai._litellm_compat import BudgetExceededError as _BudgetExceededError
diff --git a/osmosis_ai/rubric_eval.py b/osmosis_ai/rubric_eval.py
index 352bf044..2b88c24b 100644
--- a/osmosis_ai/rubric_eval.py
+++ b/osmosis_ai/rubric_eval.py
@@ -12,7 +12,6 @@
import json
import os
import re
-import warnings
from typing import Any
from ._litellm_compat import (
@@ -431,18 +430,6 @@ def _call_litellm(
# Let LiteLLM silently drop the param for models that don't support it
completion_kwargs["drop_params"] = True
- # Suppress Pydantic serialization warnings caused by LiteLLM model
- # definitions not matching provider responses. Applied as a persistent
- # filter because some providers (e.g. Gemini) trigger the warning lazily
- # after the completion call returns.
- # See: https://github.com/BerriAI/litellm/issues/17631
- warnings.filterwarnings(
- "ignore",
- message="Pydantic serializer warnings",
- category=UserWarning,
- module=r"pydantic\.main",
- )
-
try:
response = _litellm_completion(**completion_kwargs)
except _LitellmNotFoundError as err:
From 76b236c864f4c5b03dda0dcdf8ca9a12deadccee Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 12:07:54 -0800
Subject: [PATCH 58/60] Update dotenv loading in CLI to use find_dotenv with
usecwd=True for improved environment variable resolution in user project
directories.
---
osmosis_ai/cli.py | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/osmosis_ai/cli.py b/osmosis_ai/cli.py
index 85da8149..be58c7a4 100644
--- a/osmosis_ai/cli.py
+++ b/osmosis_ai/cli.py
@@ -4,7 +4,7 @@
import sys
import warnings
-from dotenv import load_dotenv
+from dotenv import find_dotenv, load_dotenv
from .cli_commands import (
EvalRubricCommand,
@@ -25,8 +25,11 @@ def main(argv: list[str] | None = None) -> int:
# Meaningful user-facing messages should use logging or print instead.
warnings.filterwarnings("ignore")
- # Load environment variables from .env file in current working directory
- load_dotenv()
+ # Load environment variables from .env file in current working directory.
+ # find_dotenv(usecwd=True) is required because the default find_dotenv()
+ # walks up from the caller's file path (site-packages or editable source),
+ # which never reaches the user's project directory.
+ load_dotenv(find_dotenv(usecwd=True))
parser = _build_parser()
args = parser.parse_args(argv)
From dbdf6f39d84988c557cc1c5a118ad12d9f5ac055 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 12:48:23 -0800
Subject: [PATCH 59/60] increment version
---
osmosis_ai/consts.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/osmosis_ai/consts.py b/osmosis_ai/consts.py
index 35f60675..1d557dab 100644
--- a/osmosis_ai/consts.py
+++ b/osmosis_ai/consts.py
@@ -1,3 +1,3 @@
# package metadata
package_name = "osmosis-ai"
-PACKAGE_VERSION = "0.2.15"
+PACKAGE_VERSION = "0.2.16"
From b01ad4ad9ada0adab0a3fb1daf092b9102b3cc51 Mon Sep 17 00:00:00 2001
From: JoyboyBrian
Date: Wed, 18 Feb 2026 13:18:12 -0800
Subject: [PATCH 60/60] [ci] chore: add permissions to workflows for content
access
---
.github/workflows/check-pr-title.yml | 3 +++
.github/workflows/tests.yml | 3 +++
2 files changed, 6 insertions(+)
diff --git a/.github/workflows/check-pr-title.yml b/.github/workflows/check-pr-title.yml
index 893c2aec..1b40f4d7 100644
--- a/.github/workflows/check-pr-title.yml
+++ b/.github/workflows/check-pr-title.yml
@@ -4,6 +4,9 @@ on:
pull_request:
types: [opened, edited, synchronize]
+permissions:
+ contents: read
+
jobs:
check-title:
runs-on: ubuntu-latest
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 545c88e3..9f00a633 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -6,6 +6,9 @@ on:
- main
pull_request:
+permissions:
+ contents: read
+
jobs:
lint:
runs-on: ubuntu-latest