diff --git a/aios/config/config.yaml.example b/aios/config/config.yaml.example index d49e05492..5dd388de0 100644 --- a/aios/config/config.yaml.example +++ b/aios/config/config.yaml.example @@ -36,6 +36,10 @@ llms: backend: "ollama" hostname: "http://localhost:11434" # Make sure to run ollama server + - name: "qwen3:4b" + backend: "ollama" + hostname: "http://localhost:11434" + # HuggingFace Models # - name: "meta-llama/Llama-3.1-8B-Instruct" # backend: "huggingface" diff --git a/aios/llm_core/adapter.py b/aios/llm_core/adapter.py index d3cca489e..3104eea00 100644 --- a/aios/llm_core/adapter.py +++ b/aios/llm_core/adapter.py @@ -17,8 +17,10 @@ import concurrent.futures from collections import defaultdict from concurrent.futures import ThreadPoolExecutor +import threading import traceback import litellm +import requests from .utils import check_availability_for_selected_llm_lists # Configure logging @@ -114,6 +116,15 @@ def __init__( self._setup_api_keys() self._initialize_llms() + # Extract Ollama hostname from the first Ollama config entry + self._ollama_hostname = "http://localhost:11434" + for llm_cfg in self.llm_configs: + if llm_cfg.backend == "ollama" and llm_cfg.hostname: + self._ollama_hostname = llm_cfg.hostname + break + + self._dynamic_registration_lock = threading.Lock() + routing_strategy = config.get_router_config().get("strategy", RouterStrategy.Sequential) # breakpoint() @@ -273,6 +284,97 @@ def _initialize_single_llm(self, config: LLMConfig) -> Optional[Union[str, HfLoc logger.error(f"Error initializing LLM {config.name} ({config.backend}): {e}", exc_info=True) return None + def _query_ollama_available_models(self) -> set: + """ + Query the Ollama server for available models. + + Returns: + A set of model name strings available on the + Ollama server, or an empty set on any failure. + """ + try: + resp = requests.get( + f"{self._ollama_hostname}/api/tags", + timeout=5 + ) + resp.raise_for_status() + data = resp.json() + models = {m["name"] for m in data.get("models", [])} + logger.debug( + f"Queried Ollama at {self._ollama_hostname}: " + f"found {len(models)} models" + ) + return models + except Exception as e: + logger.warning( + f"Failed to query Ollama models at " + f"{self._ollama_hostname}/api/tags: {e}" + ) + return set() + + def _dynamic_register_ollama_model(self, model_name: str) -> bool: + """ + Dynamically register an Ollama model that is available on the + server but not listed in config.yaml. + + Thread-safe and idempotent — calling twice for the same model + will not create duplicates. + + Args: + model_name: The Ollama model name to register. + + Returns: + True if the model is now registered (or was already), + False on any failure. + """ + try: + with self._dynamic_registration_lock: + # Re-check: another thread may have registered it + if model_name in self.available_llm_names: + return True + + available_models = self._query_ollama_available_models() + if model_name not in available_models: + logger.warning( + f"Ollama model '{model_name}' not found on " + f"server at {self._ollama_hostname}" + ) + return False + + llm_config = LLMConfig( + name=model_name, + backend="ollama", + hostname=self._ollama_hostname, + ) + initialized_model = self._initialize_single_llm( + llm_config + ) + if initialized_model is None: + logger.error( + f"Failed to initialize Ollama model " + f"'{model_name}' during dynamic registration" + ) + return False + + self.llm_configs.append(llm_config) + self.llms.append(initialized_model) + self.available_llm_names.append(model_name) + self.router.llm_configs = self.llm_configs + + logger.info( + f"Dynamically registered Ollama model " + f"'{model_name}' from server at " + f"{self._ollama_hostname}" + ) + return True + except Exception as e: + logger.error( + f"Unexpected error during dynamic registration " + f"of Ollama model '{model_name}': {e}", + exc_info=True, + ) + return False + def _handle_completion_error(self, error: Exception, model_name: Optional[str] = "Unknown") -> LLMResponse: """ Handle errors that occur during LLM completion, mapping them to LLMResponse. @@ -381,6 +483,30 @@ def execute_llm_syscalls( for i, selected_llm_list_availability in enumerate(selected_llm_lists_availability): if not selected_llm_list_availability: + # Attempt dynamic registration for unavailable Ollama models + has_ollama_unavailable = False + for llm in selected_llm_lists[i]: + if (llm["name"] not in self.available_llm_names + and llm.get("backend") == "ollama"): + has_ollama_unavailable = True + self._dynamic_register_ollama_model( + llm["name"] + ) + + if has_ollama_unavailable: + # Re-check availability after registration attempts + recheck = check_availability_for_selected_llm_lists( + self.available_llm_names, + [selected_llm_lists[i]], + ) + if recheck[0]: + executable_llm_syscalls.append(llm_syscalls[i]) + available_selected_llm_lists.append( + selected_llm_lists[i] + ) + continue + + # Reject: no Ollama models to try, or re-check failed logger.error(f"Selected LLMs are not available for syscall at index {i}") llm_syscall = llm_syscalls[i] llm_syscall.set_status("done") diff --git a/requirements.txt b/requirements.txt index 62f07aa57..67c9299b7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,6 +18,7 @@ watchdog>=2.1.9 redis>=4.5.1 sentence-transformers nltk +hypothesis scikit-learn pulp gdown diff --git a/tests/modules/llm/ollama/test_dynamic_registration_pbt.py b/tests/modules/llm/ollama/test_dynamic_registration_pbt.py new file mode 100644 index 000000000..3eef5ded5 --- /dev/null +++ b/tests/modules/llm/ollama/test_dynamic_registration_pbt.py @@ -0,0 +1,230 @@ +""" +Property-based test: Dynamic Ollama Model Registration + +Tests that _dynamic_register_ollama_model on LLMAdapter correctly +registers Ollama models available on the server but NOT listed in +config.yaml, making them pass check_availability_for_selected_llm_lists. + +The fix path is: execute_llm_syscalls -> _dynamic_register_ollama_model +-> _query_ollama_available_models -> _initialize_single_llm, which adds +the model to available_llm_names. The utility function itself remains +a pure name-in-list check. + +**Validates: Requirements 1.1, 1.2, 1.3, 1.4, 3.1-3.5** +""" + +import string +import threading +import unittest +from unittest.mock import patch, MagicMock + +from hypothesis import given, settings +from hypothesis import strategies as st + +from aios.llm_core.adapter import LLMAdapter, LLMConfig +from aios.llm_core.utils import ( + check_availability_for_selected_llm_lists, +) + +# Models registered in the adapter (simulating config.yaml) +CONFIGURED_MODELS = ["qwen3:1.7b"] + +# Models available on the Ollama server but NOT in config.yaml +SERVER_AVAILABLE_UNCONFIGURED = [ + "qwen3:4b", + "qwen3:8b", + "llama3:8b", + "mistral:7b", + "gemma2:9b", +] + +# All models the fake Ollama server reports +ALL_SERVER_MODELS = CONFIGURED_MODELS + SERVER_AVAILABLE_UNCONFIGURED + +# Strategy: pick from server-available-but-unconfigured set +ollama_unconfigured_model = st.sampled_from( + SERVER_AVAILABLE_UNCONFIGURED +) + +# Strategy: random model names guaranteed NOT to appear anywhere +random_unavailable_model = st.text( + alphabet=string.ascii_lowercase + string.digits, + min_size=3, + max_size=12, +).map(lambda s: f"zz_fake_{s}:latest") + + +def _make_ollama_tags_response(): + """Build a fake /api/tags JSON response.""" + return { + "models": [{"name": m} for m in ALL_SERVER_MODELS] + } + + +def _make_mock_adapter(): + """ + Create a minimal adapter-like object with the attributes + that _dynamic_register_ollama_model needs, without going + through the full LLMAdapter.__init__. + """ + adapter = object.__new__(LLMAdapter) + adapter.available_llm_names = list(CONFIGURED_MODELS) + adapter.llm_configs = [ + LLMConfig( + name="qwen3:1.7b", + backend="ollama", + hostname="http://localhost:11434", + ) + ] + adapter.llms = ["ollama/qwen3:1.7b"] + adapter._ollama_hostname = "http://localhost:11434" + adapter._dynamic_registration_lock = threading.Lock() + adapter.router = MagicMock() + adapter.router.llm_configs = adapter.llm_configs + return adapter + + +# --------------------------------------------------------------- +# Property 1: Bug Condition — Dynamic Registration Works +# --------------------------------------------------------------- + +class TestBugConditionExploration(unittest.TestCase): + """ + Tests the fix path: _dynamic_register_ollama_model queries + the Ollama server, registers the model, and then + check_availability_for_selected_llm_lists returns [True]. + """ + + @given(model_name=ollama_unconfigured_model) + @settings(max_examples=20) + def test_unconfigured_ollama_model_should_be_available( + self, model_name + ): + adapter = _make_mock_adapter() + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = ( + _make_ollama_tags_response() + ) + mock_resp.raise_for_status = MagicMock() + + with patch( + "aios.llm_core.adapter.requests.get", + return_value=mock_resp, + ): + registered = ( + adapter._dynamic_register_ollama_model( + model_name + ) + ) + + self.assertTrue( + registered, + f"_dynamic_register_ollama_model returned False " + f"for '{model_name}' which is on the server", + ) + self.assertIn(model_name, adapter.available_llm_names) + + result = check_availability_for_selected_llm_lists( + adapter.available_llm_names, + [[{"name": model_name, "backend": "ollama"}]], + ) + self.assertEqual(result, [True]) + + def test_concrete_bug_case_qwen3_4b(self): + adapter = _make_mock_adapter() + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = ( + _make_ollama_tags_response() + ) + mock_resp.raise_for_status = MagicMock() + + with patch( + "aios.llm_core.adapter.requests.get", + return_value=mock_resp, + ): + registered = ( + adapter._dynamic_register_ollama_model( + "qwen3:4b" + ) + ) + + self.assertTrue(registered) + result = check_availability_for_selected_llm_lists( + adapter.available_llm_names, + [[{"name": "qwen3:4b", "backend": "ollama"}]], + ) + self.assertEqual(result, [True]) + + +# --------------------------------------------------------------- +# Property 2: Preservation — Existing Behavior Unchanged +# --------------------------------------------------------------- + +class TestPreservation(unittest.TestCase): + """ + Verifies that EXISTING correct behavior of + check_availability_for_selected_llm_lists is preserved. + """ + + @given(model_name=st.sampled_from(CONFIGURED_MODELS)) + @settings(max_examples=20) + def test_configured_model_is_available(self, model_name): + result = check_availability_for_selected_llm_lists( + CONFIGURED_MODELS, + [[{"name": model_name, "backend": "ollama"}]], + ) + self.assertEqual(result, [True]) + + @given(model_name=random_unavailable_model) + @settings(max_examples=50) + def test_truly_unavailable_model_is_rejected( + self, model_name + ): + result = check_availability_for_selected_llm_lists( + CONFIGURED_MODELS, + [[{"name": model_name, "backend": "ollama"}]], + ) + self.assertEqual(result, [False]) + + def test_mixed_configured_and_unconfigured_returns_false( + self, + ): + result = check_availability_for_selected_llm_lists( + CONFIGURED_MODELS, + [ + [ + {"name": "qwen3:1.7b", "backend": "ollama"}, + { + "name": "nonexistent:latest", + "backend": "ollama", + }, + ] + ], + ) + self.assertEqual(result, [False]) + + def test_multiple_requests_all_configured(self): + lists = [ + [{"name": "qwen3:1.7b", "backend": "ollama"}], + [{"name": "qwen3:1.7b", "backend": "ollama"}], + [{"name": "qwen3:1.7b", "backend": "ollama"}], + ] + result = check_availability_for_selected_llm_lists( + CONFIGURED_MODELS, lists + ) + self.assertEqual(result, [True, True, True]) + + def test_empty_available_list_rejects_everything(self): + result = check_availability_for_selected_llm_lists( + [], + [[{"name": "qwen3:1.7b", "backend": "ollama"}]], + ) + self.assertEqual(result, [False]) + + +if __name__ == "__main__": + unittest.main()