Merge branch 'main' into feat/add-dana-agent-provider-stub

tests/unit/providers/inline/inference/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

tests/unit/providers/inline/inference/test_meta_reference.py

@@ -0,0 +1,44 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from unittest.mock import Mock
+
+import pytest
+
+from llama_stack.providers.inline.inference.meta_reference.model_parallel import (
+    ModelRunner,
+)
+
+
+class TestModelRunner:
+    """Test ModelRunner task dispatching for model-parallel inference."""
+
+    def test_chat_completion_task_dispatch(self):
+        """Verify ModelRunner correctly dispatches chat_completion tasks."""
+        # Create a mock generator
+        mock_generator = Mock()
+        mock_generator.chat_completion = Mock(return_value=iter([]))
+
+        runner = ModelRunner(mock_generator)
+
+        # Create a chat_completion task
+        fake_params = {"model": "test"}
+        fake_messages = [{"role": "user", "content": "test"}]
+        task = ("chat_completion", [fake_params, fake_messages])
+
+        # Execute task
+        runner(task)
+
+        # Verify chat_completion was called with correct arguments
+        mock_generator.chat_completion.assert_called_once_with(fake_params, fake_messages)
+
+    def test_invalid_task_type_raises_error(self):
+        """Verify ModelRunner rejects invalid task types."""
+        mock_generator = Mock()
+        runner = ModelRunner(mock_generator)
+
+        with pytest.raises(ValueError, match="Unexpected task type"):
+            runner(("invalid_task", []))