admin/OpenJarvis
1
0
Fork 0
mirror of https://github.com/open-jarvis/OpenJarvis.git synced 2026-07-29 18:40:38 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
cf29ef710e2191c74eeea80a3d84761dd01acfc6
OpenJarvis/tests/agents/test_orchestrator.py
T
Prathap PandGitHub c2756964a7 fix(channels): wire channel→agent handler and fix Telegram send pipeline (#94) 
* fix(channels): wire channel→agent handler and fix Telegram send pipeline

* format code

* add supported tests
2026-03-20 18:35:18 -07:00

668 lines
23 KiB
Python
Raw Blame History

"""Tests for the OrchestratorAgent."""
from __future__ import annotations
from unittest.mock import MagicMock
from openjarvis.agents._stubs import AgentContext
from openjarvis.agents.orchestrator import OrchestratorAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.types import Conversation, Message, Role, ToolResult
from openjarvis.tools._stubs import BaseTool, ToolSpec
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _CalculatorStub(BaseTool):
tool_id = "calculator"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="calculator",
description="Math calculator.",
parameters={
"type": "object",
"properties": {"expression": {"type": "string"}},
"required": ["expression"],
},
)
def execute(self, **params) -> ToolResult:
expr = params.get("expression", "0")
try:
val = eval(expr) # noqa: S307 — safe in tests
except Exception as e:
return ToolResult(tool_name="calculator", content=str(e), success=False)
return ToolResult(tool_name="calculator", content=str(val), success=True)
class _ThinkStub(BaseTool):
tool_id = "think"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="think",
description="Thinking tool.",
parameters={
"type": "object",
"properties": {"thought": {"type": "string"}},
},
)
def execute(self, **params) -> ToolResult:
return ToolResult(
tool_name="think",
content=params.get("thought", ""),
success=True,
)
def _make_engine_no_tools(content: str = "Final answer.") -> MagicMock:
"""Engine that never returns tool calls."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": content,
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "stop",
}
return engine
def _make_engine_with_tool_call(
tool_name: str = "calculator",
arguments: str = '{"expression":"2+2"}',
tool_call_id: str = "call_1",
final_content: str = "The answer is 4.",
) -> MagicMock:
"""Engine that returns one tool call then a final answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
# First call: tool call
{
"content": "",
"tool_calls": [
{"id": tool_call_id, "name": tool_name, "arguments": arguments}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "tool_calls",
},
# Second call: final answer
{
"content": final_content,
"usage": {"prompt_tokens": 15, "completion_tokens": 5, "total_tokens": 20},
"model": "test-model",
"finish_reason": "stop",
},
]
return engine
def _make_engine_multi_tool() -> MagicMock:
"""Engine that calls multiple tools in one turn."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
{
"content": "",
"tool_calls": [
{
"id": "call_1",
"name": "calculator",
"arguments": '{"expression":"2+2"}',
},
{
"id": "call_2",
"name": "think",
"arguments": '{"thought":"thinking..."}',
},
],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "tool_calls",
},
{
"content": "Done.",
"usage": {"prompt_tokens": 20, "completion_tokens": 3, "total_tokens": 23},
"model": "test-model",
"finish_reason": "stop",
},
]
return engine
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestOrchestratorAgent:
def test_agent_id(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model")
assert agent.agent_id == "orchestrator"
def test_no_tools_single_turn(self):
engine = _make_engine_no_tools("Hello!")
agent = OrchestratorAgent(engine, "test-model")
result = agent.run("Hello")
assert result.content == "Hello!"
assert result.turns == 1
assert result.tool_results == []
def test_single_tool_call(self):
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
result = agent.run("What is 2+2?")
assert result.content == "The answer is 4."
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].tool_name == "calculator"
assert result.tool_results[0].content == "4"
def test_multiple_tool_calls_same_turn(self):
engine = _make_engine_multi_tool()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub(), _ThinkStub()],
)
result = agent.run("Think and calculate.")
assert result.content == "Done."
assert result.turns == 2
assert len(result.tool_results) == 2
def test_with_context_conversation(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model")
conv = Conversation()
conv.add(Message(role=Role.SYSTEM, content="Be helpful."))
ctx = AgentContext(conversation=conv)
agent.run("Hi", context=ctx)
call_args = engine.generate.call_args
messages = call_args[0][0]
assert len(messages) == 2
assert messages[0].role == Role.SYSTEM
def test_tools_passed_to_engine(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
agent.run("Hello")
call_kwargs = engine.generate.call_args[1]
assert "tools" in call_kwargs
assert len(call_kwargs["tools"]) == 1
def test_no_tools_no_tools_kwarg(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model")
agent.run("Hello")
call_kwargs = engine.generate.call_args[1]
assert "tools" not in call_kwargs
def test_max_turns_exceeded(self):
"""When the engine keeps returning tool calls, stop after max_turns."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": "",
"tool_calls": [
{"id": "c1", "name": "calculator", "arguments": '{"expression":"1+1"}'}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "tool_calls",
}
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
max_turns=3,
)
result = agent.run("Loop forever")
assert result.turns == 3
assert result.metadata.get("max_turns_exceeded") is True
def test_unknown_tool_in_response(self):
engine = _make_engine_with_tool_call(
tool_name="unknown_tool",
arguments="{}",
final_content="Handled.",
)
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
result = agent.run("Use unknown tool")
assert result.content == "Handled."
assert len(result.tool_results) == 1
assert result.tool_results[0].success is False
def test_temperature_passthrough(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model", temperature=0.1)
agent.run("Hello")
call_kwargs = engine.generate.call_args[1]
assert call_kwargs["temperature"] == 0.1
def test_max_tokens_passthrough(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model", max_tokens=256)
agent.run("Hello")
call_kwargs = engine.generate.call_args[1]
assert call_kwargs["max_tokens"] == 256
def test_event_bus_agent_events(self):
bus = EventBus(record_history=True)
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model", bus=bus)
agent.run("Hello")
event_types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in event_types
assert EventType.AGENT_TURN_END in event_types
def test_event_bus_inference_events(self):
"""INFERENCE_START/END are now published by InstrumentedEngine,
not by agents directly. Agent tests verify agent-level events."""
bus = EventBus(record_history=True)
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model", bus=bus)
agent.run("Hello")
event_types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in event_types
assert EventType.AGENT_TURN_END in event_types
def test_event_bus_tool_events(self):
bus = EventBus(record_history=True)
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
bus=bus,
)
agent.run("Calc 2+2")
event_types = [e.event_type for e in bus.history]
assert EventType.TOOL_CALL_START in event_types
assert EventType.TOOL_CALL_END in event_types
def test_messages_accumulate(self):
"""After tool call, messages include assistant + tool messages."""
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
agent.run("What is 2+2?")
# Second call should include accumulated messages
second_call = engine.generate.call_args_list[1]
messages = second_call[0][0]
roles = [m.role for m in messages]
assert Role.ASSISTANT in roles
assert Role.TOOL in roles
def test_tool_message_has_tool_call_id(self):
engine = _make_engine_with_tool_call(tool_call_id="abc123")
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
agent.run("What is 2+2?")
second_call = engine.generate.call_args_list[1]
messages = second_call[0][0]
tool_msgs = [m for m in messages if m.role == Role.TOOL]
assert len(tool_msgs) == 1
assert tool_msgs[0].tool_call_id == "abc123"
def test_no_bus_works(self):
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
result = agent.run("What is 2+2?")
assert result.content == "The answer is 4."
def test_empty_tools_list(self):
engine = _make_engine_no_tools()
agent = OrchestratorAgent(engine, "test-model", tools=[])
result = agent.run("Hello")
assert result.content == "Final answer."
def test_three_turn_conversation(self):
"""Engine calls a tool twice before answering."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
{
"content": "",
"tool_calls": [
{
"id": "c1",
"name": "calculator",
"arguments": '{"expression":"2+2"}',
}
],
"usage": {
"prompt_tokens": 5,
"completion_tokens": 3,
"total_tokens": 8,
},
"model": "test-model",
"finish_reason": "tool_calls",
},
{
"content": "",
"tool_calls": [
{
"id": "c2",
"name": "calculator",
"arguments": '{"expression":"4*3"}',
}
],
"usage": {
"prompt_tokens": 15,
"completion_tokens": 3,
"total_tokens": 18,
},
"model": "test-model",
"finish_reason": "tool_calls",
},
{
"content": "2+2=4, 4*3=12",
"usage": {
"prompt_tokens": 25,
"completion_tokens": 5,
"total_tokens": 30,
},
"model": "test-model",
"finish_reason": "stop",
},
]
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
result = agent.run("Calculate")
assert result.turns == 3
assert len(result.tool_results) == 2
assert result.tool_results[0].content == "4"
assert result.tool_results[1].content == "12"
def test_tool_result_latency_tracked(self):
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
)
result = agent.run("What is 2+2?")
assert result.tool_results[0].latency_seconds >= 0
def test_max_turns_1(self):
"""With max_turns=1 and a tool call, should stop after 1 turn."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": "",
"tool_calls": [
{"id": "c1", "name": "calculator", "arguments": '{"expression":"1"}'}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "tool_calls",
}
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
max_turns=1,
)
result = agent.run("Calc")
assert result.turns == 1
assert result.metadata.get("max_turns_exceeded") is True
def test_agent_turn_end_data_no_tools(self):
bus = EventBus(record_history=True)
engine = _make_engine_no_tools("reply")
agent = OrchestratorAgent(engine, "test-model", bus=bus)
agent.run("Hi")
end = [e for e in bus.history if e.event_type == EventType.AGENT_TURN_END][0]
assert end.data["turns"] == 1
assert end.data["content_length"] == 5
def test_result_content_on_max_turns(self):
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": "partial",
"tool_calls": [
{"id": "c1", "name": "calculator", "arguments": '{"expression":"1"}'}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
"model": "test-model",
"finish_reason": "tool_calls",
}
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
max_turns=2,
)
result = agent.run("Calc")
# Should use the partial content if available
assert result.content == "partial"
class TestOrchestratorStructuredMode:
"""Tests for the structured (THOUGHT/TOOL/INPUT/FINAL_ANSWER) mode."""
def test_structured_mode_final_answer(self):
"""Structured mode should parse FINAL_ANSWER: correctly."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": "THOUGHT: Easy question.\nFINAL_ANSWER: Paris",
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
agent = OrchestratorAgent(
engine,
"test-model",
mode="structured",
)
result = agent.run("What is the capital of France?")
assert result.content == "Paris"
assert result.turns == 1
assert result.tool_results == []
def test_structured_mode_tool_call(self):
"""Parse TOOL:/INPUT:, execute tool, return final answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
{
"content": (
"THOUGHT: Need to calculate.\n"
"TOOL: calculator\n"
'INPUT: {"expression":"2+2"}'
),
"usage": {
"prompt_tokens": 10,
"completion_tokens": 10,
"total_tokens": 20,
},
"model": "test-model",
"finish_reason": "stop",
},
{
"content": ("THOUGHT: Got 4.\nFINAL_ANSWER: The answer is 4."),
"usage": {
"prompt_tokens": 20,
"completion_tokens": 10,
"total_tokens": 30,
},
"model": "test-model",
"finish_reason": "stop",
},
]
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
mode="structured",
)
result = agent.run("What is 2+2?")
assert result.content == "The answer is 4."
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].tool_name == "calculator"
assert result.tool_results[0].content == "4"
def test_structured_mode_enriched_descriptions(self):
"""Structured mode system prompt should contain enriched tool descriptions."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": "FINAL_ANSWER: ok",
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
mode="structured",
)
agent.run("Hello")
call_args = engine.generate.call_args
messages = call_args[0][0]
system_msg = messages[0].content
assert "### calculator" in system_msg
assert "expression" in system_msg
class TestOrchestratorParallelTools:
"""Tests for parallel tool execution."""
def test_parallel_tool_execution(self):
"""Multiple tool calls execute in parallel and return in correct order."""
import time
class _SlowTool(BaseTool):
tool_id = "slow"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="slow",
description="Slow tool.",
parameters={
"type": "object",
"properties": {"id": {"type": "string"}},
},
)
def execute(self, **params) -> ToolResult:
time.sleep(0.1) # Simulate slow operation
return ToolResult(
tool_name="slow",
content=f"result_{params.get('id', '')}",
success=True,
)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
{
"content": "",
"tool_calls": [
{"id": "c1", "name": "slow", "arguments": '{"id":"1"}'},
{"id": "c2", "name": "slow", "arguments": '{"id":"2"}'},
{"id": "c3", "name": "slow", "arguments": '{"id":"3"}'},
],
"usage": {
"prompt_tokens": 5,
"completion_tokens": 3,
"total_tokens": 8,
},
"model": "test-model",
"finish_reason": "tool_calls",
},
{
"content": "All done.",
"usage": {
"prompt_tokens": 20,
"completion_tokens": 3,
"total_tokens": 23,
},
"model": "test-model",
"finish_reason": "stop",
},
]
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_SlowTool()],
parallel_tools=True,
)
t0 = time.time()
result = agent.run("Do things")
elapsed = time.time() - t0
assert result.content == "All done."
assert len(result.tool_results) == 3
# Results should be in original order
assert result.tool_results[0].content == "result_1"
assert result.tool_results[1].content == "result_2"
assert result.tool_results[2].content == "result_3"
# Should be parallel — 3 tools at 0.1s each should take < 0.25s, not 0.3s+
assert elapsed < 0.25
def test_sequential_tool_execution(self):
"""parallel_tools=False runs tools sequentially."""
engine = _make_engine_multi_tool()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub(), _ThinkStub()],
parallel_tools=False,
)
result = agent.run("Do things")
assert result.content == "Done."
assert len(result.tool_results) == 2
def test_single_tool_call_no_parallel(self):
"""Single tool call should not use parallel path even if parallel_tools=True."""
engine = _make_engine_with_tool_call()
agent = OrchestratorAgent(
engine,
"test-model",
tools=[_CalculatorStub()],
parallel_tools=True,
)
result = agent.run("What is 2+2?")
assert result.content == "The answer is 4."
Reference in New Issue View Git Blame Copy Permalink