diff --git a/src/timmy/agents/base.py b/src/timmy/agents/base.py index b74c4248..ed8e8792 100644 --- a/src/timmy/agents/base.py +++ b/src/timmy/agents/base.py @@ -14,6 +14,7 @@ import logging from abc import ABC, abstractmethod from typing import Any +import httpx from agno.agent import Agent from agno.models.ollama import Ollama @@ -120,8 +121,15 @@ class BaseAgent(ABC): Returns: Agent response """ - result = self.agent.run(message, stream=False) - response = result.content if hasattr(result, "content") else str(result) + try: + result = self.agent.run(message, stream=False) + response = result.content if hasattr(result, "content") else str(result) + except (httpx.ConnectError, httpx.ReadError, ConnectionError) as exc: + logger.error("Ollama disconnected: %s", exc) + raise + except Exception as exc: + logger.error("Agent run failed: %s", exc) + raise # Emit completion event if self.event_bus: diff --git a/src/timmy/session.py b/src/timmy/session.py index f0ae99d2..13b80f3d 100644 --- a/src/timmy/session.py +++ b/src/timmy/session.py @@ -11,6 +11,8 @@ let Agno's session_id mechanism handle conversation continuity. import logging import re +import httpx + logger = logging.getLogger(__name__) # Default session ID for the dashboard (stable across requests) @@ -83,6 +85,9 @@ async def chat(message: str, session_id: str | None = None) -> str: try: run = await agent.arun(message, stream=False, session_id=sid) response_text = run.content if hasattr(run, "content") else str(run) + except (httpx.ConnectError, httpx.ReadError, ConnectionError) as exc: + logger.error("Ollama disconnected: %s", exc) + return "Ollama appears to be disconnected. Check that ollama serve is running." except Exception as exc: logger.error("Session: agent.arun() failed: %s", exc) return "I'm having trouble reaching my language model right now. Please try again shortly." @@ -111,6 +116,11 @@ async def chat_with_tools(message: str, session_id: str | None = None): try: return await agent.arun(message, stream=False, session_id=sid) + except (httpx.ConnectError, httpx.ReadError, ConnectionError) as exc: + logger.error("Ollama disconnected: %s", exc) + return _ErrorRunOutput( + "Ollama appears to be disconnected. Check that ollama serve is running." + ) except Exception as exc: logger.error("Session: agent.arun() failed: %s", exc) # Return a duck-typed object that callers can handle uniformly @@ -133,6 +143,11 @@ async def continue_chat(run_output, session_id: str | None = None): try: return await agent.acontinue_run(run_response=run_output, stream=False, session_id=sid) + except (httpx.ConnectError, httpx.ReadError, ConnectionError) as exc: + logger.error("Ollama disconnected: %s", exc) + return _ErrorRunOutput( + "Ollama appears to be disconnected. Check that ollama serve is running." + ) except Exception as exc: logger.error("Session: agent.acontinue_run() failed: %s", exc) return _ErrorRunOutput(f"Error continuing run: {exc}") diff --git a/tests/timmy/test_ollama_disconnect.py b/tests/timmy/test_ollama_disconnect.py new file mode 100644 index 00000000..c0ccd14e --- /dev/null +++ b/tests/timmy/test_ollama_disconnect.py @@ -0,0 +1,304 @@ +"""Test Ollama disconnection handling. + +Verifies that: +1. BaseAgent.run() logs 'Ollama disconnected' when agent.run() raises connection errors +2. BaseAgent.run() re-raises the error (not silently swallowed) +3. session.chat() returns disconnect-specific message on connection errors +4. session.chat_with_tools() returns _ErrorRunOutput with disconnect message on connection errors +""" + +import importlib +import logging +from unittest.mock import MagicMock, patch + +import httpx +import pytest + + +class TestBaseAgentDisconnect: + """Test BaseAgent.run() disconnection handling.""" + + def test_base_agent_logs_on_connect_error(self, caplog): + """BaseAgent.run() logs 'Ollama disconnected' on httpx.ConnectError.""" + caplog.set_level(logging.ERROR) + importlib.import_module("timmy.agents.base") + + with ( + patch("timmy.agents.base.Ollama") as mock_ollama, + patch("timmy.agents.base.Agent") as mock_agent_class, + ): + mock_ollama.return_value = MagicMock() + mock_agent = MagicMock() + mock_agent.run.side_effect = httpx.ConnectError("Connection refused") + mock_agent_class.return_value = mock_agent + + from timmy.agents.base import BaseAgent + + class ConcreteAgent(BaseAgent): + async def execute_task(self, task_id: str, description: str, context: dict): + return {"task_id": task_id, "status": "completed"} + + agent = ConcreteAgent( + agent_id="test", + name="Test", + role="tester", + system_prompt="You are a test agent.", + tools=[], + ) + + with pytest.raises(httpx.ConnectError): + import asyncio + + asyncio.run(agent.run("test message")) + + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + def test_base_agent_logs_on_read_error(self, caplog): + """BaseAgent.run() logs 'Ollama disconnected' on httpx.ReadError.""" + caplog.set_level(logging.ERROR) + importlib.import_module("timmy.agents.base") + + with ( + patch("timmy.agents.base.Ollama") as mock_ollama, + patch("timmy.agents.base.Agent") as mock_agent_class, + ): + mock_ollama.return_value = MagicMock() + mock_agent = MagicMock() + mock_agent.run.side_effect = httpx.ReadError("Server closed connection") + mock_agent_class.return_value = mock_agent + + from timmy.agents.base import BaseAgent + + class ConcreteAgent(BaseAgent): + async def execute_task(self, task_id: str, description: str, context: dict): + return {"task_id": task_id, "status": "completed"} + + agent = ConcreteAgent( + agent_id="test", + name="Test", + role="tester", + system_prompt="You are a test agent.", + tools=[], + ) + + with pytest.raises(httpx.ReadError): + import asyncio + + asyncio.run(agent.run("test message")) + + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + def test_base_agent_logs_on_connection_error(self, caplog): + """BaseAgent.run() logs 'Ollama disconnected' on ConnectionError.""" + caplog.set_level(logging.ERROR) + importlib.import_module("timmy.agents.base") + + with ( + patch("timmy.agents.base.Ollama") as mock_ollama, + patch("timmy.agents.base.Agent") as mock_agent_class, + ): + mock_ollama.return_value = MagicMock() + mock_agent = MagicMock() + mock_agent.run.side_effect = ConnectionError("Network unreachable") + mock_agent_class.return_value = mock_agent + + from timmy.agents.base import BaseAgent + + class ConcreteAgent(BaseAgent): + async def execute_task(self, task_id: str, description: str, context: dict): + return {"task_id": task_id, "status": "completed"} + + agent = ConcreteAgent( + agent_id="test", + name="Test", + role="tester", + system_prompt="You are a test agent.", + tools=[], + ) + + with pytest.raises(ConnectionError): + import asyncio + + asyncio.run(agent.run("test message")) + + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + def test_base_agent_re_raises_connection_error(self): + """BaseAgent.run() re-raises the connection error (not silently swallowed).""" + importlib.import_module("timmy.agents.base") + + with ( + patch("timmy.agents.base.Ollama") as mock_ollama, + patch("timmy.agents.base.Agent") as mock_agent_class, + ): + mock_ollama.return_value = MagicMock() + mock_agent = MagicMock() + mock_agent.run.side_effect = httpx.ConnectError("Connection refused") + mock_agent_class.return_value = mock_agent + + from timmy.agents.base import BaseAgent + + class ConcreteAgent(BaseAgent): + async def execute_task(self, task_id: str, description: str, context: dict): + return {"task_id": task_id, "status": "completed"} + + agent = ConcreteAgent( + agent_id="test", + name="Test", + role="tester", + system_prompt="You are a test agent.", + tools=[], + ) + + with pytest.raises(httpx.ConnectError, match="Connection refused"): + import asyncio + + asyncio.run(agent.run("test message")) + + +class TestSessionDisconnect: + """Test session.py disconnection handling.""" + + @pytest.mark.asyncio + async def test_chat_returns_disconnect_message_on_connect_error(self, caplog): + """session.chat() returns disconnect-specific message on httpx.ConnectError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = httpx.ConnectError("Connection refused") + mock_get_agent.return_value = mock_agent + + # Import after patching + from timmy import session + + result = await session.chat("test message") + + assert "Ollama appears to be disconnected" in result + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_chat_returns_disconnect_message_on_read_error(self, caplog): + """session.chat() returns disconnect-specific message on httpx.ReadError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = httpx.ReadError("Server closed connection") + mock_get_agent.return_value = mock_agent + + from timmy import session + + result = await session.chat("test message") + + assert "Ollama appears to be disconnected" in result + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_chat_returns_disconnect_message_on_connection_error(self, caplog): + """session.chat() returns disconnect-specific message on ConnectionError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = ConnectionError("Network unreachable") + mock_get_agent.return_value = mock_agent + + from timmy import session + + result = await session.chat("test message") + + assert "Ollama appears to be disconnected" in result + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_chat_with_tools_returns_error_run_output_on_connect_error(self, caplog): + """session.chat_with_tools() returns _ErrorRunOutput with disconnect message on ConnectError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = httpx.ConnectError("Connection refused") + mock_get_agent.return_value = mock_agent + + from timmy import session + + result = await session.chat_with_tools("test message") + + assert hasattr(result, "content") + assert hasattr(result, "status") + assert "Ollama appears to be disconnected" in result.content + assert result.status == "ERROR" + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_chat_with_tools_returns_error_run_output_on_read_error(self, caplog): + """session.chat_with_tools() returns _ErrorRunOutput with disconnect message on ReadError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = httpx.ReadError("Server closed connection") + mock_get_agent.return_value = mock_agent + + from timmy import session + + result = await session.chat_with_tools("test message") + + assert "Ollama appears to be disconnected" in result.content + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_continue_chat_returns_error_run_output_on_connect_error(self, caplog): + """session.continue_chat() returns _ErrorRunOutput with disconnect message on ConnectError.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.acontinue_run.side_effect = httpx.ConnectError("Connection refused") + mock_get_agent.return_value = mock_agent + + from timmy import session + + mock_run_output = MagicMock() + result = await session.continue_chat(mock_run_output) + + assert hasattr(result, "content") + assert "Ollama appears to be disconnected" in result.content + assert any("Ollama disconnected" in record.message for record in caplog.records), ( + f"Expected 'Ollama disconnected' in logs, got: {[r.message for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_other_errors_use_generic_message(self, caplog): + """Non-connection errors still use the generic error message.""" + caplog.set_level(logging.ERROR) + + with patch("timmy.session._get_agent") as mock_get_agent: + mock_agent = MagicMock() + mock_agent.arun.side_effect = ValueError("Some other error") + mock_get_agent.return_value = mock_agent + + from timmy import session + + result = await session.chat("test message") + + assert "I'm having trouble reaching my language model" in result + # Should NOT have Ollama disconnected message + assert "Ollama appears to be disconnected" not in result