feat(ai_client): Emit API lifecycle and tool execution events
This commit is contained in:
@@ -18,3 +18,97 @@ def test_event_emission():
|
||||
ai_client.events.emit("request_start", payload={"model": "test"})
|
||||
|
||||
mock_callback.assert_called_once_with(payload={"model": "test"})
|
||||
|
||||
def test_send_emits_events():
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# We need to mock _ensure_gemini_client and the chat object it creates
|
||||
with patch("ai_client._ensure_gemini_client"), \
|
||||
patch("ai_client._gemini_client") as mock_client, \
|
||||
patch("ai_client._gemini_chat") as mock_chat:
|
||||
|
||||
# Setup mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.candidates = []
|
||||
# Explicitly set usage_metadata as a mock with integer values
|
||||
mock_usage = MagicMock()
|
||||
mock_usage.prompt_token_count = 10
|
||||
mock_usage.candidates_token_count = 5
|
||||
mock_usage.cached_content_token_count = None
|
||||
mock_response.usage_metadata = mock_usage
|
||||
mock_chat.send_message.return_value = mock_response
|
||||
mock_client.chats.create.return_value = mock_chat
|
||||
|
||||
ai_client.set_provider("gemini", "gemini-flash")
|
||||
|
||||
start_callback = MagicMock()
|
||||
response_callback = MagicMock()
|
||||
|
||||
ai_client.events.on("request_start", start_callback)
|
||||
ai_client.events.on("response_received", response_callback)
|
||||
|
||||
# We need to bypass the context changed check or set it up
|
||||
ai_client.send("context", "message")
|
||||
|
||||
assert start_callback.called
|
||||
assert response_callback.called
|
||||
|
||||
# Check payload
|
||||
args, kwargs = start_callback.call_args
|
||||
assert kwargs['payload']['provider'] == 'gemini'
|
||||
|
||||
def test_send_emits_tool_events():
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
with patch("ai_client._ensure_gemini_client"), \
|
||||
patch("ai_client._gemini_client") as mock_client, \
|
||||
patch("ai_client._gemini_chat") as mock_chat, \
|
||||
patch("mcp_client.dispatch") as mock_dispatch:
|
||||
|
||||
# 1. Setup mock response with a tool call
|
||||
mock_fc = MagicMock()
|
||||
mock_fc.name = "read_file"
|
||||
mock_fc.args = {"path": "test.txt"}
|
||||
|
||||
mock_response_with_tool = MagicMock()
|
||||
mock_response_with_tool.candidates = [MagicMock()]
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = "tool call text"
|
||||
mock_part.function_call = mock_fc
|
||||
mock_response_with_tool.candidates[0].content.parts = [mock_part]
|
||||
mock_response_with_tool.candidates[0].finish_reason.name = "STOP"
|
||||
|
||||
# Setup mock usage
|
||||
mock_usage = MagicMock()
|
||||
mock_usage.prompt_token_count = 10
|
||||
mock_usage.candidates_token_count = 5
|
||||
mock_usage.cached_content_token_count = None
|
||||
mock_response_with_tool.usage_metadata = mock_usage
|
||||
|
||||
# 2. Setup second mock response (final answer)
|
||||
mock_response_final = MagicMock()
|
||||
mock_response_final.candidates = []
|
||||
mock_response_final.usage_metadata = mock_usage
|
||||
|
||||
mock_chat.send_message.side_effect = [mock_response_with_tool, mock_response_final]
|
||||
mock_dispatch.return_value = "file content"
|
||||
|
||||
ai_client.set_provider("gemini", "gemini-flash")
|
||||
|
||||
tool_callback = MagicMock()
|
||||
ai_client.events.on("tool_execution", tool_callback)
|
||||
|
||||
ai_client.send("context", "message")
|
||||
|
||||
# Should be called twice: once for 'started', once for 'completed'
|
||||
assert tool_callback.call_count == 2
|
||||
|
||||
# Check 'started' call
|
||||
args, kwargs = tool_callback.call_args_list[0]
|
||||
assert kwargs['payload']['status'] == 'started'
|
||||
assert kwargs['payload']['tool'] == 'read_file'
|
||||
|
||||
# Check 'completed' call
|
||||
args, kwargs = tool_callback.call_args_list[1]
|
||||
assert kwargs['payload']['status'] == 'completed'
|
||||
assert kwargs['payload']['result'] == 'file content'
|
||||
|
||||
Reference in New Issue
Block a user