Fix tests

This commit is contained in:
Cyber MacGeddon 2025-11-28 16:36:29 +00:00
parent e24de6081f
commit a9a1da115c
2 changed files with 6 additions and 6 deletions

View file

@ -182,8 +182,8 @@ Final Answer: Machine learning is a field of AI that enables computers to learn
assert action.observation == "Machine learning is a subset of AI that enables computers to learn from data."
# Verify callbacks were called
think_callback.assert_called_once_with("I need to search for information about machine learning")
observe_callback.assert_called_once_with("Machine learning is a subset of AI that enables computers to learn from data.")
think_callback.assert_called_once_with("I need to search for information about machine learning", is_final=True)
observe_callback.assert_called_once_with("Machine learning is a subset of AI that enables computers to learn from data.", is_final=True)
# Verify tool was executed
graph_rag_client = mock_flow_context("graph-rag-request")
@ -211,7 +211,7 @@ Final Answer: Machine learning is a branch of artificial intelligence."""
assert action.final == "Machine learning is a branch of artificial intelligence."
# Verify only think callback was called (no observation for final answer)
think_callback.assert_called_once_with("I can provide a direct answer")
think_callback.assert_called_once_with("I can provide a direct answer", is_final=True)
observe_callback.assert_not_called()
@pytest.mark.asyncio
@ -457,7 +457,7 @@ Args: {args_json}"""
# Assert
assert isinstance(action, Action)
assert action.observation == expected_response.strip()
observe_callback.assert_called_with(expected_response.strip())
observe_callback.assert_called_with(expected_response.strip(), is_final=True)
# Reset mocks
mock_flow_context("graph-rag-request").reset_mock()

View file

@ -346,7 +346,7 @@ class AgentManager:
# In non-streaming mode, send complete thought
# In streaming mode, thoughts were already sent as chunks
if not streaming:
await think(act.thought)
await think(act.thought, is_final=True)
return act
else:
@ -354,7 +354,7 @@ class AgentManager:
# In non-streaming mode, send complete thought
# In streaming mode, thoughts were already sent as chunks
if not streaming:
await think(act.thought)
await think(act.thought, is_final=True)
logger.debug(f"ACTION: {act.name}")