68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
from beaver.engine.providers.base import LLMProvider, LLMResponse
|
|
from beaver.engine.providers.factory import ProviderBundle
|
|
from beaver.tasks import TaskExecutionPlanner, TaskRecord
|
|
|
|
|
|
class _TeamPlannerProvider(LLMProvider):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.calls = 0
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict],
|
|
tools: list[dict] | None = None,
|
|
model: str | None = None,
|
|
max_tokens: int = 4096,
|
|
temperature: float = 0.7,
|
|
) -> LLMResponse:
|
|
self.calls += 1
|
|
return LLMResponse(
|
|
content='{"mode":"team","reason":"parallel research","strategy":"parallel","nodes":[{"node_id":"research","task":"research","agent":{"name":"researcher"}}]}',
|
|
finish_reason="stop",
|
|
provider_name="stub",
|
|
model="stub-model",
|
|
)
|
|
|
|
def get_default_model(self) -> str:
|
|
return "stub-model"
|
|
|
|
|
|
def test_agent_team_can_be_disabled_by_environment(monkeypatch) -> None:
|
|
monkeypatch.setenv("BEAVER_AGENT_TEAM_ENABLED", "0")
|
|
provider = _TeamPlannerProvider()
|
|
task = TaskRecord(
|
|
task_id="task-1",
|
|
session_id="session-1",
|
|
description="research and compare options",
|
|
goal="research and compare options",
|
|
constraints=[],
|
|
priority=0,
|
|
status="open",
|
|
creator="test",
|
|
created_at="now",
|
|
updated_at="now",
|
|
)
|
|
bundle = ProviderBundle(
|
|
main_runtime=SimpleNamespace(model="stub-model", provider_name="stub"),
|
|
main_provider=provider,
|
|
)
|
|
|
|
plan = asyncio.run(
|
|
TaskExecutionPlanner().plan(
|
|
task=task,
|
|
user_message="research and compare options",
|
|
attempt_index=1,
|
|
provider_bundle=bundle,
|
|
)
|
|
)
|
|
|
|
assert plan.mode == "single"
|
|
assert plan.reason == "planner_disabled_by_environment"
|
|
assert provider.calls == 0
|