d33db9fff8
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Clément Siriex <clement.sirieix@mistral.ai> Co-Authored-By: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai> Co-Authored-By: Thaddee Tyl <thaddee.tyl@gmail.com> Co-Authored-By: David Brochart <david.brochart@gmail.com> Co-Authored-By: Joseph Guhlin <joseph.guhlin@gmail.com> Co-Authored-By: Thomas Kenbeek <thomaskenbeek@gmail.com> Co-Authored-By: Remenby31 <baptiste.cruvellier31@gmail.com>
134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from vibe.core.agents.models import BUILTIN_AGENTS, AgentProfile, BuiltinAgentName
|
|
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
|
from vibe.core.middleware import (
|
|
PLAN_AGENT_REMINDER,
|
|
ConversationContext,
|
|
MiddlewareAction,
|
|
MiddlewarePipeline,
|
|
PlanAgentMiddleware,
|
|
)
|
|
from vibe.core.types import AgentStats
|
|
|
|
|
|
def make_context() -> ConversationContext:
|
|
config = VibeConfig(session_logging=SessionLoggingConfig(enabled=False))
|
|
return ConversationContext(messages=[], stats=AgentStats(), config=config)
|
|
|
|
|
|
class TestPlanAgentMiddleware:
|
|
@pytest.mark.asyncio
|
|
async def test_injects_reminder_when_plan_agent_active(self) -> None:
|
|
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
|
assert result.message == PLAN_AGENT_REMINDER
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_does_not_inject_when_default_agent(self) -> None:
|
|
middleware = PlanAgentMiddleware(
|
|
lambda: BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
|
)
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
assert result.message is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_does_not_inject_when_auto_approve_agent(self) -> None:
|
|
middleware = PlanAgentMiddleware(
|
|
lambda: BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
|
)
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
assert result.message is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_does_not_inject_when_accept_edits_agent(self) -> None:
|
|
middleware = PlanAgentMiddleware(
|
|
lambda: BUILTIN_AGENTS[BuiltinAgentName.ACCEPT_EDITS]
|
|
)
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
assert result.message is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_after_turn_always_continues(self) -> None:
|
|
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
|
ctx = make_context()
|
|
|
|
result = await middleware.after_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dynamically_checks_agent(self) -> None:
|
|
current_profile: AgentProfile = BUILTIN_AGENTS[BuiltinAgentName.DEFAULT]
|
|
middleware = PlanAgentMiddleware(lambda: current_profile)
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
|
|
current_profile = BUILTIN_AGENTS[BuiltinAgentName.PLAN]
|
|
result = await middleware.before_turn(ctx)
|
|
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
|
|
|
current_profile = BUILTIN_AGENTS[BuiltinAgentName.AUTO_APPROVE]
|
|
result = await middleware.before_turn(ctx)
|
|
assert result.action == MiddlewareAction.CONTINUE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_reminder(self) -> None:
|
|
custom_reminder = "Custom plan agent reminder"
|
|
middleware = PlanAgentMiddleware(
|
|
lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN], reminder=custom_reminder
|
|
)
|
|
ctx = make_context()
|
|
|
|
result = await middleware.before_turn(ctx)
|
|
|
|
assert result.message == custom_reminder
|
|
|
|
def test_reset_does_nothing(self) -> None:
|
|
middleware = PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN])
|
|
middleware.reset()
|
|
|
|
|
|
class TestMiddlewarePipelineWithPlanAgent:
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_includes_plan_agent_injection(self) -> None:
|
|
pipeline = MiddlewarePipeline()
|
|
pipeline.add(PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.PLAN]))
|
|
ctx = make_context()
|
|
|
|
result = await pipeline.run_before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.INJECT_MESSAGE
|
|
assert PLAN_AGENT_REMINDER in (result.message or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_skips_injection_when_not_plan_agent(self) -> None:
|
|
pipeline = MiddlewarePipeline()
|
|
pipeline.add(
|
|
PlanAgentMiddleware(lambda: BUILTIN_AGENTS[BuiltinAgentName.DEFAULT])
|
|
)
|
|
ctx = make_context()
|
|
|
|
result = await pipeline.run_before_turn(ctx)
|
|
|
|
assert result.action == MiddlewareAction.CONTINUE
|