placeholder/features/steps/stream_router_agent_tool_steps.py

"""Steps for SimpleToolAgent and SimpleLLMAgent coverage."""

from __future__ import annotations

from typing import Any

from behave import given, then, when
from behave.runner import Context

import cleveragents.reactive.stream_router as stream_router
from cleveragents.core.exceptions import StreamRoutingError
from cleveragents.reactive.stream_router import SimpleLLMAgent, SimpleToolAgent


class _WorkingTemplate:
    def __init__(self, template: str) -> None:
        self.template = template

    def render(self, **kwargs: Any) -> str:
        name = kwargs.get("name")
        if name is not None:
            return self.template.replace("{{ name }}", str(name))
        return f"rendered:{self.template}"


class _WorkingTemplateEnv:
    def from_string(self, template: str) -> _WorkingTemplate:
        return _WorkingTemplate(template)


class _FailingTemplate:
    def render(self, **_kwargs: Any) -> str:
        raise RuntimeError("render failed")


class _FailingTemplateEnv:
    def from_string(self, _template: str) -> _FailingTemplate:
        return _FailingTemplate()


class _MessageStub:
    def __init__(self, content: str) -> None:
        self.content = content


class _LLMStub:
    def __init__(self, response_content: str = "llm-output") -> None:
        self.response_content = response_content
        self.invocations: list[list[Any]] = []

    def invoke(self, messages: list[Any]) -> Any:
        self.invocations.append(messages)

        class _Response:
            def __init__(self, content: str) -> None:
                self.content = content

        return _Response(self.response_content)


class _RegistryStub:
    def __init__(self, llm: _LLMStub) -> None:
        self.llm = llm
        self.calls: list[dict[str, Any]] = []

    def create_llm(self, **kwargs: Any) -> _LLMStub:
        self.calls.append(kwargs)
        return self.llm


def _register_cleanup(context: Context, cleanup_fn) -> None:
    cleanup = getattr(context, "add_cleanup", None)
    if callable(cleanup):
        cleanup(cleanup_fn)


@given("a SimpleToolAgent with a tool missing code")
def step_tool_agent_missing_code(context: Context) -> None:
    context.tool_agent = SimpleToolAgent(tools=[{}])


@given("a SimpleToolAgent with a code block tool")
def step_tool_agent_code_block(context: Context) -> None:
    context.tool_agent = SimpleToolAgent(
        tools=[{"code": "result = 'ok:' + str(input_data)"}]
    )


@when('I process "{content}" through the SimpleToolAgent')
def step_process_tool_agent(context: Context, content: str) -> None:
    context.tool_result = context.tool_agent.process(content)


@when('I process "{content}" through the SimpleToolAgent expecting rejection')
def step_process_tool_agent_reject(context: Context, content: str) -> None:
    context.tool_error = None
    try:
        context.tool_result = context.tool_agent.process(content)
    except Exception as exc:  # pylint: disable=broad-except
        context.tool_error = exc


@when('I call process_message_sync with "{content}"')
def step_process_tool_agent_sync(context: Context, content: str) -> None:
    context.tool_result = context.tool_agent.process_message_sync(content)


@when('I call process_message_sync with "{content}" expecting rejection')
def step_process_tool_agent_sync_reject(context: Context, content: str) -> None:
    context.tool_error = None
    try:
        context.tool_result = context.tool_agent.process_message_sync(content)
    except Exception as exc:  # pylint: disable=broad-except
        context.tool_error = exc


@then('the SimpleToolAgent result should be "{expected}"')
def step_tool_agent_result(context: Context, expected: str) -> None:
    assert context.tool_result == expected


@then("the SimpleToolAgent result should be empty")
def step_tool_agent_result_empty(context: Context) -> None:
    assert context.tool_result == ""


@then("the SimpleToolAgent should raise a routing error about code blocks")
def step_tool_agent_code_block_rejected(context: Context) -> None:
    assert isinstance(context.tool_error, StreamRoutingError), (
        f"Expected StreamRoutingError, got {type(context.tool_error).__name__}: "
        f"{context.tool_error}"
    )
    assert "code" in str(context.tool_error).lower()
    assert "not supported" in str(context.tool_error).lower()


@given("a SimpleLLMAgent for prompt rendering")
def step_llm_agent_for_rendering(context: Context) -> None:
    context.llm_agent = SimpleLLMAgent("renderer", {"system_prompt": ""})


@given('the SimpleLLMAgent template environment is "{mode}"')
def step_set_template_env(context: Context, mode: str) -> None:
    if mode == "none":
        context.llm_agent._template_env = None
        return
    if mode == "working":
        context.llm_agent._template_env = _WorkingTemplateEnv()
        return
    if mode == "failing":
        context.llm_agent._template_env = _FailingTemplateEnv()
        return
    raise ValueError(f"Unknown template env mode: {mode}")


@when('I render the template "{template}" with context:')
def step_render_template(context: Context, template: str) -> None:
    table = getattr(context, "table", None)
    payload = {row["key"]: row["value"] for row in table} if table else {}
    if template == "<empty>":
        template = ""
    context.rendered_prompt = context.llm_agent._render_prompt(template, payload)


@then('the rendered prompt should be "{expected}"')
def step_rendered_prompt(context: Context, expected: str) -> None:
    if expected == "<empty>":
        expected = ""
    assert context.rendered_prompt == expected


@given("a SimpleLLMAgent with provider config")
def step_llm_agent_with_provider_config(context: Context) -> None:
    context.llm_agent = SimpleLLMAgent(
        "agent",
        {
            "provider": "stub",
            "model": "stub-model",
            "temperature": 0.7,
            "max_tokens": 50,
            "max_retries": 2,
        },
    )


@given("a SimpleLLMAgent with system prompt and provider config")
def step_llm_agent_with_system_prompt(context: Context) -> None:
    context.llm_agent = SimpleLLMAgent(
        "agent",
        {
            "provider": "stub",
            "model": "stub-model",
            "temperature": 0.7,
            "max_tokens": 50,
            "max_retries": 2,
            "system_prompt": "system",
        },
    )


@given("a stub provider registry is installed")
def step_stub_provider_registry(context: Context) -> None:
    context.llm_stub = _LLMStub()
    context.registry_stub = _RegistryStub(context.llm_stub)
    context.original_registry = stream_router.get_provider_registry
    stream_router.get_provider_registry = lambda: context.registry_stub
    _register_cleanup(
        context,
        lambda: setattr(
            stream_router, "get_provider_registry", context.original_registry
        ),
    )


@given("LangChain message classes are unavailable")
def step_langchain_unavailable(context: Context) -> None:
    context.original_human = stream_router.HumanMessage
    context.original_system = stream_router.SystemMessage
    stream_router.HumanMessage = None
    stream_router.SystemMessage = None
    _register_cleanup(
        context,
        lambda: (
            setattr(stream_router, "HumanMessage", context.original_human),
            setattr(stream_router, "SystemMessage", context.original_system),
        ),
    )


@given("LangChain message stubs are installed")
def step_langchain_stubbed(context: Context) -> None:
    context.original_human = stream_router.HumanMessage
    context.original_system = stream_router.SystemMessage
    stream_router.HumanMessage = _MessageStub
    stream_router.SystemMessage = _MessageStub
    _register_cleanup(
        context,
        lambda: (
            setattr(stream_router, "HumanMessage", context.original_human),
            setattr(stream_router, "SystemMessage", context.original_system),
        ),
    )


@when("I resolve the LLM twice")
def step_resolve_llm_twice(context: Context) -> None:
    context.llm_first = context.llm_agent._resolve_llm()
    context.llm_second = context.llm_agent._resolve_llm()


@then("the registry should be called once with the configured kwargs")
def step_verify_registry_calls(context: Context) -> None:
    assert len(context.registry_stub.calls) == 1
    call = context.registry_stub.calls[0]
    assert call.get("provider_type") == "stub"
    assert call.get("model_id") == "stub-model"
    assert call.get("temperature") == 0.7
    assert call.get("max_tokens") == 50
    assert call.get("max_retries") == 2


@then("the resolved LLM should be cached")
def step_verify_llm_cached(context: Context) -> None:
    assert context.llm_first is context.llm_second


@when('I process "{content}" through the LLM agent')
def step_process_llm_agent(context: Context, content: str) -> None:
    context.process_error = None
    try:
        context.process_result = context.llm_agent.process(content)
    except Exception as exc:  # pylint: disable=broad-except
        context.process_error = exc


@then("I should receive a stream routing error about missing messages")
def step_verify_missing_messages_error(context: Context) -> None:
    assert isinstance(context.process_error, StreamRoutingError)
    assert "LangChain messages not available" in str(context.process_error)


@when("I process 123 through the LLM agent via process_message_sync")
def step_process_llm_agent_sync(context: Context) -> None:
    context.process_result = context.llm_agent.process_message_sync(123)


@then("the LLM should receive system and human messages")
def step_verify_llm_messages(context: Context) -> None:
    assert context.llm_stub.invocations
    messages = context.llm_stub.invocations[0]
    assert len(messages) == 2
    assert isinstance(messages[0], _MessageStub)
    assert isinstance(messages[1], _MessageStub)
    # System prompt is augmented with boundary instructions (mechanism 2)
    assert "system" in messages[0].content
    # User content is wrapped with boundary markers (mechanism 2)
    assert "123" in messages[1].content


@then('the LLM agent result should be "llm-output"')
def step_verify_llm_result(context: Context) -> None:
    assert context.process_result == "llm-output"