.

Author: basnijholt

adaptive_scheduler/_server_support/llm_manager.py

@@ -1,20 +1,26 @@
 from __future__ import annotations
 
 import asyncio
-import random
+from typing import TYPE_CHECKING
 
 import aiofiles
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import HumanMessage, SystemMessage
 
 from .base_manager import BaseManager
 
+if TYPE_CHECKING:
+    from langchain.schema import BaseMessage
+
 
 class LLMManager(BaseManager):
     """A manager for handling interactions with a language model."""
 
-    def __init__(self) -> None:
+    def __init__(self, model_name: str = "gpt-4") -> None:
         super().__init__()
+        self.llm = ChatOpenAI(model_name=model_name)
         self._diagnoses_cache: dict[str, str] = {}
-        self._chat_history: list[dict[str, str]] = []
+        self._chat_history: list[BaseMessage] = []
 
     async def _manage(self) -> None:
         """The main loop for the manager."""
@@ -47,31 +53,24 @@ async def diagnose_failed_job(self, job_id: str) -> str:
         log_content = await self._read_log_file(log_path)
         if log_content == "Log file not found.":
             return log_content
-        # In a real implementation, this would be an API call to an LLM
-        diagnosis = await self._simulate_llm_call(
-            f"Analyze the following log and determine the cause of failure:\n\n{log_content}",
-        )
+
+        messages: list[BaseMessage] = [
+            SystemMessage(
+                content="You are a helpful assistant that analyzes job failure logs.",
+            ),
+            HumanMessage(
+                content=f"Analyze the following log and determine the cause of failure:\n\n{log_content}",
+            ),
+        ]
+        response = await self.llm.agenerate([messages])
+        diagnosis = response.generations[0][0].text
         self._diagnoses_cache[job_id] = diagnosis
         return diagnosis
 
     async def chat(self, message: str) -> str:
         """Handles a chat message and returns a response."""
-        self._chat_history.append({"role": "user", "content": message})
-        # In a real implementation, this would be an API call to an LLM
-        response = await self._simulate_llm_call(str(self._chat_history))
-        self._chat_history.append({"role": "assistant", "content": response})
-        return response
-
-    async def _simulate_llm_call(self, prompt: str) -> str:  # noqa: ARG002
-        """Simulates a call to a language model."""
-        await asyncio.sleep(
-            random.uniform(0.1, 0.5),  # noqa: S311
-        )  # Simulate network latency
-        responses = [
-            "It seems like there was a `FileNotFoundError`. Check if the input files are correctly specified.",
-            "The job failed due to a `MemoryError`. Try requesting more memory for your job.",
-            "I see a `ValueError` in the logs. It seems like an invalid argument was passed to a function.",
-            "The simulation diverged. You might want to adjust the simulation parameters.",
-            "I'm not sure what went wrong. Could you provide more details?",
-        ]
-        return random.choice(responses)  # noqa: S311
+        self._chat_history.append(HumanMessage(content=message))
+        response = await self.llm.agenerate([self._chat_history])
+        result = response.generations[0][0].text
+        self._chat_history.append(SystemMessage(content=result))
+        return result

adaptive_scheduler/widgets.py

@@ -959,9 +959,12 @@ async def on_submit(sender: ipyw.Text) -> None:
         sender.value = ""
         if run_manager.llm_manager is None:
             return
-        response = await run_manager.llm_manager.chat(message)
-        chat_history.value += f"You: {message}\n"
-        chat_history.value += f"LLM: {response}\n"
+        try:
+            response = await run_manager.llm_manager.chat(message)
+            chat_history.value += f"You: {message}\n"
+            chat_history.value += f"LLM: {response}\n"
+        except Exception as e:  # noqa: BLE001
+            chat_history.value += f"Error: {e}\n"
 
     def on_submit_wrapper(sender: ipyw.Text) -> None:
         task = asyncio.create_task(on_submit(sender))
@@ -986,8 +989,11 @@ async def on_failed_job_change(change: dict[str, Any]) -> None:
         job_id = change["new"]
         if run_manager.llm_manager is None:
             return
-        diagnosis = await run_manager.llm_manager.diagnose_failed_job(job_id)
-        chat_history.value = f"Diagnosis for job {job_id}:\n{diagnosis}\n"
+        try:
+            diagnosis = await run_manager.llm_manager.diagnose_failed_job(job_id)
+            chat_history.value = f"Diagnosis for job {job_id}:\n{diagnosis}\n"
+        except Exception as e:  # noqa: BLE001
+            chat_history.value = f"Error: {e}\n"
 
     def on_failed_job_change_wrapper(change: dict[str, Any]) -> None:
         task = asyncio.create_task(on_failed_job_change(change))

tests/test_llm_manager.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
 from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -17,35 +16,54 @@
 
 
 @pytest.mark.asyncio
-async def test_diagnose_failed_job() -> None:
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+async def test_diagnose_failed_job(mock_chat_openai: MagicMock) -> None:
     """Test that the diagnose_failed_job method returns a diagnosis."""
+    mock_llm = mock_chat_openai.return_value
+    mock_llm.agenerate = AsyncMock(
+        return_value=MagicMock(
+            generations=[[MagicMock(text="diagnosis")]],
+        ),
+    )
+
     llm_manager = LLMManager()
     job_id = "test_job"
-    with patch("adaptive_scheduler._server_support.llm_manager.aiofiles.open") as mock_open:
+    with patch(
+        "adaptive_scheduler._server_support.llm_manager.aiofiles.open",
+    ) as mock_open:
         mock_open.return_value.__aenter__.return_value.read.return_value = (
             "This is a log file with an error."
         )
         diagnosis = await llm_manager.diagnose_failed_job(job_id)
-    assert isinstance(diagnosis, str)
+    assert diagnosis == "diagnosis"
 
 
 @pytest.mark.asyncio
-async def test_chat() -> None:
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+async def test_chat(mock_chat_openai: MagicMock) -> None:
     """Test that the chat method returns a response."""
+    mock_llm = mock_chat_openai.return_value
+    mock_llm.agenerate = AsyncMock(
+        return_value=MagicMock(
+            generations=[[MagicMock(text="response")]],
+        ),
+    )
     llm_manager = LLMManager()
     message = "Hello, world!"
     response = await llm_manager.chat(message)
-    assert isinstance(response, str)
+    assert response == "response"
 
 
 @pytest.fixture
-def llm_manager() -> LLMManager:
-    """An LLMManager instance."""
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+def llm_manager(mock_chat_openai: MagicMock) -> LLMManager:  # noqa: ARG001
+    """An LLMManager instance with a mocked ChatOpenAI."""
     return LLMManager()
 
 
 @pytest.fixture
-def run_manager() -> RunManager:
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+def run_manager(mock_chat_openai: MagicMock) -> RunManager:  # noqa: ARG001
     """A RunManager instance with a mocked scheduler."""
     scheduler = MagicMock()
     learners = [MagicMock()]
@@ -103,25 +121,26 @@ async def test_job_manager_diagnoses_failed_job_async(
 async def test_llm_manager_cache(llm_manager: LLMManager) -> None:
     """Test that the LLMManager caches diagnoses."""
     job_id = "test_job"
-    with (
-        patch.object(
-            llm_manager,
-            "_read_log_file",
-            return_value="log content",
+    llm_manager.llm.agenerate = AsyncMock(
+        return_value=MagicMock(
+            generations=[[MagicMock(text="diagnosis")]],
         ),
-        patch.object(
-            llm_manager,
-            "_simulate_llm_call",
-            return_value="diagnosis",
-        ) as mock_llm_call,
+    )
+    with patch.object(
+        llm_manager,
+        "_read_log_file",
+        return_value="log content",
     ):
         await llm_manager.diagnose_failed_job(job_id)
         await llm_manager.diagnose_failed_job(job_id)
-        mock_llm_call.assert_called_once()
+        llm_manager.llm.agenerate.assert_called_once()
 
 
 @pytest.mark.asyncio
-async def test_diagnose_failed_job_file_not_found() -> None:
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+async def test_diagnose_failed_job_file_not_found(
+    mock_chat_openai: MagicMock,  # noqa: ARG001
+) -> None:
     """Test that the diagnose_failed_job method handles a missing log file."""
     llm_manager = LLMManager()
     job_id = "test_job"
@@ -136,13 +155,22 @@ async def test_diagnose_failed_job_file_not_found() -> None:
 @pytest.mark.asyncio
 async def test_chat_history(llm_manager: LLMManager) -> None:
     """Test that the chat history is maintained."""
+    llm_manager.llm.agenerate = AsyncMock(
+        return_value=MagicMock(
+            generations=[[MagicMock(text="response")]],
+        ),
+    )
     await llm_manager.chat("Hello")
     await llm_manager.chat("How are you?")
     assert len(llm_manager._chat_history) == 4
 
 
 @pytest.mark.asyncio
-async def test_read_log_file(tmp_path: Path) -> None:
+@patch("adaptive_scheduler._server_support.llm_manager.ChatOpenAI")
+async def test_read_log_file(
+    mock_chat_openai: MagicMock,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
     """Test that the _read_log_file method reads a file asynchronously."""
     llm_manager = LLMManager()
     log_content = "This is a test log file."
@@ -151,70 +179,3 @@ async def test_read_log_file(tmp_path: Path) -> None:
 
     read_content = await llm_manager._read_log_file(str(log_file))
     assert read_content == log_content
-
-
-@pytest.mark.asyncio
-async def test_chat_widget_callbacks(run_manager: RunManager) -> None:
-    """Test the async callbacks of the chat_widget."""
-    from adaptive_scheduler.widgets import chat_widget
-
-    # Mock the chat_widget dependencies
-    with (
-        patch("ipywidgets.Text") as mock_text,
-        patch("ipywidgets.Textarea") as mock_textarea,
-        patch("ipywidgets.Dropdown") as mock_dropdown,
-        patch("ipywidgets.VBox"),
-        patch("adaptive_scheduler.widgets._add_title"),
-    ):
-        # Create instances of the mocked widgets
-        text_input = mock_text.return_value
-        chat_history = mock_textarea.return_value
-        failed_job_dropdown = mock_dropdown.return_value
-
-        # Call the widget function to set up the callbacks
-        chat_widget(run_manager)
-
-        # --- Test on_submit ---
-        # Get the on_submit wrapper from the mock
-        on_submit_wrapper = text_input.on_submit.call_args[0][0]
-
-        # Get a reference to the value mock before it's changed
-        chat_history_value_mock = chat_history.value
-        chat_history_value_mock.__iadd__.return_value = chat_history_value_mock
-
-        # Mock the chat method
-        with patch.object(
-            run_manager.llm_manager,
-            "chat",
-            return_value="Test response",
-        ) as mock_chat:
-            # Simulate the submission
-            text_input.value = "Test message"
-            on_submit_wrapper(text_input)
-            await asyncio.sleep(0)  # allow the task to run
-
-            # Assertions
-            mock_chat.assert_called_once_with("Test message")
-            iadd_calls = chat_history_value_mock.__iadd__.call_args_list
-            assert "You: Test message" in iadd_calls[0].args[0]
-            assert "LLM: Test response" in iadd_calls[1].args[0]
-
-        # --- Test on_failed_job_change ---
-        # Get the on_failed_job_change wrapper from the mock
-        on_failed_job_change_wrapper = failed_job_dropdown.observe.call_args[0][0]
-
-        # Mock the diagnose_failed_job method
-        with patch.object(
-            run_manager.llm_manager,
-            "diagnose_failed_job",
-            return_value="Test diagnosis",
-        ) as mock_diagnose:
-            # Simulate the dropdown change
-            change = {"new": "test_job_id"}
-            on_failed_job_change_wrapper(change)
-            await asyncio.sleep(0)  # allow the task to run
-
-            # Assertions
-            mock_diagnose.assert_called_once_with("test_job_id")
-            assert "Diagnosis for job test_job_id" in chat_history.value
-            assert "Test diagnosis" in chat_history.value