diff --git a/examples/01_standalone_sdk/37_llm_profile_store.py b/examples/01_standalone_sdk/37_llm_profile_store.py new file mode 100644 index 0000000000..ee484e8fa6 --- /dev/null +++ b/examples/01_standalone_sdk/37_llm_profile_store.py @@ -0,0 +1,72 @@ +"""Example: Using LLMProfileStore to save and reuse LLM configurations. + +LLMProfileStore persists LLM configurations as JSON files, so you can define +a profile once and reload it across sessions without repeating setup code. +""" + +import os +import tempfile + +from pydantic import SecretStr + +from openhands.sdk import LLM, LLMProfileStore + + +# Use a temporary directory so this example doesn't pollute your home folder. +# In real usage you can omit base_dir to use the default (~/.openhands/profiles). +store = LLMProfileStore(base_dir=tempfile.mkdtemp()) + + +# 1. Create two LLM profiles with different usage + +api_key = os.getenv("LLM_API_KEY") +assert api_key is not None, "LLM_API_KEY environment variable is not set." +base_url = os.getenv("LLM_BASE_URL") +model = os.getenv("LLM_MODEL", "anthropic/claude-sonnet-4-5-20250929") + +fast_llm = LLM( + usage_id="fast", + model=model, + api_key=SecretStr(api_key), + base_url=base_url, + temperature=0.0, +) + +creative_llm = LLM( + usage_id="creative", + model=model, + api_key=SecretStr(api_key), + base_url=base_url, + temperature=0.9, +) + +# 2. Save profiles + +# Note that secrets are excluded by default for safety. +store.save("fast", fast_llm) +store.save("creative", creative_llm) + +# To persist the API key as well, pass `include_secrets=True`: +# store.save("fast", fast_llm, include_secrets=True) + +# 3. List available persisted profiles + +print(f"Stored profiles: {store.list()}") + +# 4. Load a profile + +loaded = store.load("fast") +assert isinstance(loaded, LLM) +print( + "Loaded profile. " + f"usage:{loaded.usage_id}, " + f"model: {loaded.model}, " + f"temperature: {loaded.temperature}." +) + +# 5. Delete a profile + +store.delete("creative") +print(f"After deletion: {store.list()}") + +print("EXAMPLE_COST: 0") diff --git a/openhands-sdk/openhands/sdk/__init__.py b/openhands-sdk/openhands/sdk/__init__.py index 1ea20a995f..811b2ac47c 100644 --- a/openhands-sdk/openhands/sdk/__init__.py +++ b/openhands-sdk/openhands/sdk/__init__.py @@ -25,6 +25,7 @@ from openhands.sdk.llm import ( LLM, ImageContent, + LLMProfileStore, LLMRegistry, LLMStreamChunk, Message, @@ -66,6 +67,7 @@ __all__ = [ "LLM", "LLMRegistry", + "LLMProfileStore", "LLMStreamChunk", "TokenCallbackType", "ConversationStats", diff --git a/openhands-sdk/openhands/sdk/llm/__init__.py b/openhands-sdk/openhands/sdk/llm/__init__.py index b2b5ad3047..7431bdd50e 100644 --- a/openhands-sdk/openhands/sdk/llm/__init__.py +++ b/openhands-sdk/openhands/sdk/llm/__init__.py @@ -5,6 +5,7 @@ OpenAISubscriptionAuth, ) from openhands.sdk.llm.llm import LLM +from openhands.sdk.llm.llm_profile_store import LLMProfileStore from openhands.sdk.llm.llm_registry import LLMRegistry, RegistryEvent from openhands.sdk.llm.llm_response import LLMResponse from openhands.sdk.llm.message import ( @@ -37,6 +38,7 @@ "LLMResponse", "LLM", "LLMRegistry", + "LLMProfileStore", "RouterLLM", "RegistryEvent", # Messages diff --git a/openhands-sdk/openhands/sdk/llm/llm_profile_store.py b/openhands-sdk/openhands/sdk/llm/llm_profile_store.py new file mode 100644 index 0000000000..6dbd5695d2 --- /dev/null +++ b/openhands-sdk/openhands/sdk/llm/llm_profile_store.py @@ -0,0 +1,180 @@ +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING, Final + +from filelock import FileLock, Timeout + +from openhands.sdk.logger import get_logger + + +if TYPE_CHECKING: + from openhands.sdk.llm.llm import LLM + +_DEFAULT_PROFILE_DIR: Final[Path] = Path.home() / ".openhands" / "profiles" +_LOCK_TIMEOUT_SECONDS: Final[float] = 30.0 + +logger = get_logger(__name__) + + +class LLMProfileStore: + """Standalone utility for persisting LLM configurations.""" + + def __init__(self, base_dir: Path | str | None = None) -> None: + """Initialize the profile store. + + Args: + base_dir: Path to the directory where the profiles are stored. + If `None` is provided, the default directory is used, i.e., + `~/.openhands/profiles`. + """ + self.base_dir = Path(base_dir) if base_dir is not None else _DEFAULT_PROFILE_DIR + # ensure directory existence + self.base_dir.mkdir(parents=True, exist_ok=True) + self._file_lock = FileLock(self.base_dir / ".profiles.lock") + + @contextmanager + def _acquire_lock(self, timeout: float = _LOCK_TIMEOUT_SECONDS) -> Iterator[None]: + """Acquire file lock for safe concurrent access. + + Args: + timeout: Maximum time to wait for lock acquisition in seconds. + + Raises: + TimeoutError: If the lock cannot be acquired within the timeout. + """ + try: + with self._file_lock.acquire(timeout=timeout): + yield + except Timeout: + logger.error(f"[Profile Store] Failed to acquire lock within {timeout}s") + raise TimeoutError( + f"Profile store lock acquisition timed out after {timeout}s" + ) + + def list(self) -> list[str]: + """Returns a list of all profiles stored. + + Returns: + List of profile filenames (e.g., ["default.json", "gpt4.json"]). + """ + with self._acquire_lock(): + return [p.name for p in self.base_dir.glob("*.json")] + + def _get_profile_path(self, name: str) -> Path: + """Get the full path for a profile name. + + Args: + name: Profile name (must be a simple filename without path separators) + + Raises: + ValueError: If name contains path separators or is invalid + """ + # Remove .json extension if present for consistent handling + clean_name = name.removesuffix(".json") + + # Validate: no path separators, not empty, no hidden files + if ( + not clean_name + or "/" in clean_name + or "\\" in clean_name + or clean_name.startswith(".") + ): + raise ValueError( + f"Invalid profile name: {name!r}. " + "Profile names must be simple filenames without path separators." + ) + + return self.base_dir / f"{clean_name}.json" + + def save(self, name: str, llm: "LLM", include_secrets: bool = False) -> None: + """Save a profile to the profile directory. + + Note that if a profile name already exists, it will be overwritten. + + Args: + name: Name of the profile to save. + llm: LLM instance to save + include_secrets: Whether to include the profile secrets. Defaults to False. + + Raises: + TimeoutError: If the lock cannot be acquired. + """ + profile_path = self._get_profile_path(name) + + with self._acquire_lock(): + if profile_path.exists(): + logger.info( + f"[Profile Store] Profile `{name}` already exists. Overwriting." + ) + + profile_json = llm.model_dump_json( + exclude_none=True, + indent=2, + context={"expose_secrets": include_secrets}, + ) + with tempfile.NamedTemporaryFile( + mode="w", dir=self.base_dir, suffix=".tmp", delete=False + ) as tmp: + tmp.write(profile_json) + tmp_path = Path(tmp.name) + + Path.replace(tmp_path, profile_path) + logger.info(f"[Profile Store] Saved profile `{name}` at {profile_path}") + + def load(self, name: str) -> "LLM": + """Load an LLM instance from the given profile name. + + Args: + name: Name of the profile to load. + + Returns: + An LLM instance constructed from the profile configuration. + + Raises: + FileNotFoundError: If the profile name does not exist. + ValueError: If the profile file is corrupted or invalid. + TimeoutError: If the lock cannot be acquired. + """ + profile_path = self._get_profile_path(name) + + with self._acquire_lock(): + if not profile_path.exists(): + existing = [p.name for p in self.base_dir.glob("*.json")] + raise FileNotFoundError( + f"Profile `{name}` not found. " + f"Available profiles: {', '.join(existing) or 'none'}" + ) + + try: + from openhands.sdk.llm.llm import LLM + + llm_instance = LLM.load_from_json(str(profile_path)) + except Exception as e: + # Re-raise as ValueError for clearer error handling + raise ValueError(f"Failed to load profile `{name}`: {e}") from e + + logger.info(f"[Profile Store] Loaded profile `{name}` from {profile_path}") + return llm_instance + + def delete(self, name: str) -> None: + """Delete an existing profile. + + If the profile is not present in the profile directory, it does nothing. + + Args: + name: Name of the profile to delete. + + Raises: + TimeoutError: If the lock cannot be acquired. + """ + profile_path = self._get_profile_path(name) + + with self._acquire_lock(): + if not profile_path.exists(): + logger.info(f"[Profile Store] Profile `{name}` not found. Skipping.") + return + + profile_path.unlink() + logger.info(f"[Profile Store] Deleted profile `{name}`") diff --git a/tests/sdk/llm/test_llm_profile_store.py b/tests/sdk/llm/test_llm_profile_store.py new file mode 100644 index 0000000000..cc2f1fe343 --- /dev/null +++ b/tests/sdk/llm/test_llm_profile_store.py @@ -0,0 +1,406 @@ +import concurrent.futures +import json +import threading +from pathlib import Path + +import pytest +from pydantic import SecretStr + +from openhands.sdk.llm import LLM +from openhands.sdk.llm.llm_profile_store import LLMProfileStore + + +@pytest.fixture +def profile_store(tmp_path: Path) -> LLMProfileStore: + """Create a profile store with a temporary directory.""" + return LLMProfileStore(base_dir=tmp_path) + + +@pytest.fixture +def sample_llm() -> LLM: + """Create a sample LLM instance for testing.""" + return LLM( + usage_id="test-llm", + model="gpt-4", + temperature=0.7, + max_output_tokens=2000, + ) + + +@pytest.fixture +def sample_llm_with_secrets() -> LLM: + """Create a sample LLM instance with secrets for testing.""" + return LLM( + usage_id="test-llm-secrets", + model="gpt-4", + temperature=0.5, + api_key=SecretStr("secret-api-key-12345"), + ) + + +def test_init_creates_directory(tmp_path: Path) -> None: + """Test that initialization creates the base directory.""" + profile_dir = tmp_path / "profiles" + assert not profile_dir.exists() + + LLMProfileStore(base_dir=profile_dir) + + assert profile_dir.exists() + assert profile_dir.is_dir() + + +def test_init_with_string_path(tmp_path: Path) -> None: + """Test initialization with a string path.""" + profile_dir = str(tmp_path / "profiles") + store = LLMProfileStore(base_dir=profile_dir) + + assert store.base_dir == Path(profile_dir) + assert store.base_dir.exists() + + +def test_init_with_path_object(tmp_path: Path) -> None: + """Test initialization with a Path object.""" + profile_dir = tmp_path / "profiles" + store = LLMProfileStore(base_dir=profile_dir) + + assert store.base_dir == profile_dir + assert store.base_dir.exists() + + +def test_init_with_existing_directory(tmp_path: Path) -> None: + """Test initialization with an existing directory.""" + profile_dir = tmp_path / "profiles" + profile_dir.mkdir() + + store = LLMProfileStore(base_dir=profile_dir) + + assert store.base_dir == profile_dir + + +def test_list_empty_store(profile_store: LLMProfileStore) -> None: + """Test listing profiles in an empty store.""" + profiles = profile_store.list() + assert profiles == [] + + +def test_list_with_profiles(profile_store: LLMProfileStore, sample_llm: LLM) -> None: + """Test listing profiles after saving some.""" + profile_store.save("profile1", sample_llm) + profile_store.save("profile2", sample_llm) + + profiles = profile_store.list() + + assert len(profiles) == 2 + assert "profile1.json" in profiles + assert "profile2.json" in profiles + + +def test_list_excludes_non_json_files( + profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test that list() only returns .json files.""" + profile_store.save("valid", sample_llm) + + # Create a non-json file + (profile_store.base_dir / "not_a_profile.txt").write_text("hello") + + profiles = profile_store.list() + + assert profiles == ["valid.json"] + + +def test_save_creates_file(profile_store: LLMProfileStore, sample_llm: LLM) -> None: + """Test that save creates a profile file.""" + profile_store.save("my_profile", sample_llm) + + profile_path = profile_store.base_dir / "my_profile.json" + assert profile_path.exists() + + +@pytest.mark.parametrize( + "name", + ["", ".json", ".", "..", "my/profile", "my//profile"], +) +def test_save_with_invalid_profile_name( + name: str, profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + with pytest.raises(ValueError, match=f"Invalid profile name: {name!r}. "): + profile_store.save(name, sample_llm) + + +def test_save_writes_valid_json( + profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test that saved file contains valid JSON.""" + profile_store.save("my_profile", sample_llm) + + profile_path = profile_store.base_dir / "my_profile.json" + content = profile_path.read_text() + data = json.loads(content) + + assert data["model"] == "gpt-4" + assert data["temperature"] == 0.7 + + +def test_save_with_json_extension( + profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test saving with .json extension in name.""" + profile_store.save("my_profile.json", sample_llm) + + # Should not create my_profile.json.json + assert (profile_store.base_dir / "my_profile.json").exists() + assert not (profile_store.base_dir / "my_profile.json.json").exists() + + +def test_save_overwrites_existing( + profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test that save overwrites an existing profile.""" + profile_store.save("my_profile", sample_llm) + + # Modify and save again + modified_llm = LLM( + usage_id="modified", + model="gpt-3.5-turbo", + temperature=0.3, + ) + profile_store.save("my_profile", modified_llm) + + # Load and verify + loaded = profile_store.load("my_profile") + assert loaded.model == "gpt-3.5-turbo" + assert loaded.temperature == 0.3 + + +def test_save_without_secrets( + profile_store: LLMProfileStore, sample_llm_with_secrets: LLM +) -> None: + """Test that secrets are not saved by default.""" + profile_store.save("with_secrets", sample_llm_with_secrets) + + profile_path = profile_store.base_dir / "with_secrets.json" + content = profile_path.read_text() + + # Secret should be masked + assert "secret-api-key-12345" not in content + + +def test_save_with_secrets( + profile_store: LLMProfileStore, sample_llm_with_secrets: LLM +) -> None: + """Test that secrets are saved when include_secrets=True.""" + profile_store.save("with_secrets", sample_llm_with_secrets, include_secrets=True) + + profile_path = profile_store.base_dir / "with_secrets.json" + content = profile_path.read_text() + + # Secret should be present + assert "secret-api-key-12345" in content + + +@pytest.mark.parametrize("name", ["my_profile", "my_profile.json"]) +def test_load_existing_profile( + name: str, profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test loading an existing profile.""" + profile_store.save(name, sample_llm) + + loaded = profile_store.load(name) + + assert loaded.usage_id == sample_llm.usage_id + assert loaded.model == sample_llm.model + assert loaded.temperature == sample_llm.temperature + assert loaded.max_output_tokens == sample_llm.max_output_tokens + + +def test_load_nonexistent_profile(profile_store: LLMProfileStore) -> None: + """Test loading a profile that doesn't exist.""" + with pytest.raises(FileNotFoundError) as exc_info: + profile_store.load("nonexistent") + + assert "nonexistent" in str(exc_info.value) + assert "not found" in str(exc_info.value) + + +def test_load_nonexistent_shows_available( + profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test that error message shows available profiles.""" + profile_store.save("available1", sample_llm) + profile_store.save("available2", sample_llm) + + with pytest.raises(FileNotFoundError) as exc_info: + profile_store.load("nonexistent") + + error_msg = str(exc_info.value) + assert "available1.json" in error_msg + assert "available2.json" in error_msg + + +def test_load_corrupted_profile(profile_store: LLMProfileStore) -> None: + """Test loading a corrupted profile raises ValueError.""" + # Create a corrupted profile file + profile_path = profile_store.base_dir / "corrupted.json" + profile_path.write_text("{ invalid json }") + + with pytest.raises(ValueError) as exc_info: + profile_store.load("corrupted") + + assert "Failed to load profile" in str(exc_info.value) + assert "corrupted" in str(exc_info.value) + + +@pytest.mark.parametrize("name", ["to_delete", "to_delete.json"]) +def test_delete_existing_profile( + name: str, profile_store: LLMProfileStore, sample_llm: LLM +) -> None: + """Test deleting an existing profile.""" + profile_store.save(name, sample_llm) + profile_filename = f"{name}.json" if not name.endswith(".json") else name + assert profile_filename in profile_store.list() + + profile_store.delete(name) + assert profile_filename not in profile_store.list() + + +def test_delete_nonexistent_profile(profile_store: LLMProfileStore) -> None: + """Test that deleting a nonexistent profile doesn't raise an error.""" + profile_store.delete("nonexistent") + + +def test_concurrent_saves(tmp_path: Path) -> None: + """Test that concurrent saves don't corrupt data.""" + store = LLMProfileStore(base_dir=tmp_path) + num_threads = 10 + results: list[int] = [] + errors: list[tuple[int, Exception]] = [] + + def save_profile(index: int) -> None: + try: + llm = LLM( + usage_id=f"test-{index}", + model=f"model-{index}", + temperature=0.1 * index, + ) + store.save(f"profile_{index}", llm) + results.append(index) + except Exception as e: + errors.append((index, e)) + + threads = [ + threading.Thread(target=save_profile, args=(i,)) for i in range(num_threads) + ] + + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0, f"Errors occurred: {errors}" + assert len(results) == num_threads + + # Verify all profiles were saved correctly + profiles = store.list() + assert len(profiles) == num_threads + + +def test_concurrent_reads_and_writes(tmp_path: Path) -> None: + """Test concurrent reads and writes don't cause issues.""" + store = LLMProfileStore(base_dir=tmp_path) + + # Pre-create some profiles + for i in range(5): + llm = LLM(usage_id=f"test-{i}", model=f"model-{i}") + store.save(f"profile_{i}", llm) + + errors: list[tuple[str, str | int, Exception]] = [] + read_results: list[str] = [] + write_results: list[int] = [] + + def read_profile(name: str) -> None: + try: + loaded = store.load(name) + read_results.append(loaded.model) + except Exception as e: + errors.append(("read", name, e)) + + def write_profile(index: int) -> None: + try: + llm = LLM(usage_id=f"new-{index}", model=f"new-model-{index}") + store.save(f"new_profile_{index}", llm) + write_results.append(index) + except Exception as e: + errors.append(("write", index, e)) + + with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor: + futures = [] + # Submit read tasks + for i in range(5): + futures.append(executor.submit(read_profile, f"profile_{i}")) + # Submit write tasks + for i in range(5): + futures.append(executor.submit(write_profile, i)) + + concurrent.futures.wait(futures) + + assert len(errors) == 0, f"Errors occurred: {errors}" + assert len(read_results) == 5 + assert len(write_results) == 5 + + +def test_full_workflow(profile_store: LLMProfileStore) -> None: + """Test a complete save-list-load-delete workflow.""" + llm = LLM( + usage_id="workflow-test", + model="claude-3-opus", + temperature=0.8, + max_output_tokens=4096, + ) + + # Save + profile_store.save("workflow_profile", llm) + + # List + profiles = profile_store.list() + assert "workflow_profile.json" in profiles + + # Load + loaded = profile_store.load("workflow_profile") + assert loaded.usage_id == llm.usage_id + assert loaded.model == llm.model + assert loaded.temperature == llm.temperature + assert loaded.max_output_tokens == llm.max_output_tokens + + # Delete + profile_store.delete("workflow_profile") + assert "workflow_profile.json" not in profile_store.list() + + +def test_multiple_profiles(profile_store: LLMProfileStore) -> None: + """Test managing multiple profiles.""" + profiles_data = [ + ("gpt4", "gpt-4", 0.7), + ("gpt35", "gpt-3.5-turbo", 0.5), + ("claude", "claude-3-opus", 0.9), + ] + + # Save all + for name, model, temp in profiles_data: + llm = LLM(usage_id=name, model=model, temperature=temp) + profile_store.save(name, llm) + + # Verify all exist + stored = profile_store.list() + assert len(stored) == 3 + + # Load and verify each + for name, expected_model, expected_temp in profiles_data: + loaded = profile_store.load(name) + assert loaded.model == expected_model + assert loaded.temperature == expected_temp + + # Delete one + profile_store.delete("gpt35") + assert len(profile_store.list()) == 2 + assert "gpt35.json" not in profile_store.list()