forked from Rockachopa/Timmy-time-dashboard
fix: log Ollama disconnections with specific error handling (#92)
- BaseAgent.run(): catch httpx.ConnectError/ReadError/ConnectionError, log 'Ollama disconnected: <error>' at ERROR level, then re-raise - session.py: distinguish Ollama disconnects from other errors in chat(), chat_with_tools(), continue_chat() — return specific message 'Ollama appears to be disconnected' instead of generic error - 11 new tests covering all disconnect paths
This commit is contained in:
@@ -14,6 +14,7 @@ import logging
|
|||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
from agno.agent import Agent
|
from agno.agent import Agent
|
||||||
from agno.models.ollama import Ollama
|
from agno.models.ollama import Ollama
|
||||||
|
|
||||||
@@ -120,8 +121,15 @@ class BaseAgent(ABC):
|
|||||||
Returns:
|
Returns:
|
||||||
Agent response
|
Agent response
|
||||||
"""
|
"""
|
||||||
result = self.agent.run(message, stream=False)
|
try:
|
||||||
response = result.content if hasattr(result, "content") else str(result)
|
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
|
# Emit completion event
|
||||||
if self.event_bus:
|
if self.event_bus:
|
||||||
|
|||||||
@@ -11,6 +11,8 @@ let Agno's session_id mechanism handle conversation continuity.
|
|||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# Default session ID for the dashboard (stable across requests)
|
# 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:
|
try:
|
||||||
run = await agent.arun(message, stream=False, session_id=sid)
|
run = await agent.arun(message, stream=False, session_id=sid)
|
||||||
response_text = run.content if hasattr(run, "content") else str(run)
|
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:
|
except Exception as exc:
|
||||||
logger.error("Session: agent.arun() failed: %s", exc)
|
logger.error("Session: agent.arun() failed: %s", exc)
|
||||||
return "I'm having trouble reaching my language model right now. Please try again shortly."
|
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:
|
try:
|
||||||
return await agent.arun(message, stream=False, session_id=sid)
|
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:
|
except Exception as exc:
|
||||||
logger.error("Session: agent.arun() failed: %s", exc)
|
logger.error("Session: agent.arun() failed: %s", exc)
|
||||||
# Return a duck-typed object that callers can handle uniformly
|
# 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:
|
try:
|
||||||
return await agent.acontinue_run(run_response=run_output, stream=False, session_id=sid)
|
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:
|
except Exception as exc:
|
||||||
logger.error("Session: agent.acontinue_run() failed: %s", exc)
|
logger.error("Session: agent.acontinue_run() failed: %s", exc)
|
||||||
return _ErrorRunOutput(f"Error continuing run: {exc}")
|
return _ErrorRunOutput(f"Error continuing run: {exc}")
|
||||||
|
|||||||
304
tests/timmy/test_ollama_disconnect.py
Normal file
304
tests/timmy/test_ollama_disconnect.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user