cleveragents-core/features/steps/base_agent_module_steps.py

"""Behave steps targeting cleveragents.agents.base coverage."""

from __future__ import annotations

import asyncio
import importlib.util
import sys
from pathlib import Path
from types import ModuleType
from typing import Any

from behave import given, then, when

BASE_MODULE_PATH = (
    Path(__file__).resolve().parents[2] / "src" / "cleveragents" / "agents" / "base.py"
)


class RecordingApp:
    """Minimal callable graph application that records interactions."""

    def __init__(self) -> None:
        self.invoke_calls: list[dict[str, Any]] = []
        self.ainvoke_calls: list[dict[str, Any]] = []
        self.stream_calls: list[dict[str, Any]] = []

    def invoke(
        self, input_data: dict[str, Any], config: dict[str, Any]
    ) -> dict[str, Any]:
        call = {"input": input_data, "config": config}
        self.invoke_calls.append(call)
        return {"input": input_data, "config": config}

    async def ainvoke(
        self, input_data: dict[str, Any], config: dict[str, Any]
    ) -> dict[str, Any]:
        call = {"input": input_data, "config": config}
        self.ainvoke_calls.append(call)
        return {"input": input_data, "config": config, "async": True}

    def stream(self, input_data: dict[str, Any], config: dict[str, Any]):
        call = {"input": input_data, "config": config}
        self.stream_calls.append(call)
        yield {"event": "start", "config": config}
        yield {"event": "payload", "input": input_data, "config": config}
        yield {"event": "done", "config": config}


class TestGraph:
    """Graph stub that remembers the checkpointer used during compilation."""

    def __init__(self, context: Any) -> None:
        self.context = context
        self.compile_checkpointers: list[Any] = []

    def compile(self, checkpointer: Any = None) -> RecordingApp:
        self.compile_checkpointers.append(checkpointer)
        compiled_app = RecordingApp()
        self.context.recording_app = compiled_app
        return compiled_app


def _register_cleanup(context: Any, callback) -> None:
    if hasattr(context, "_cleanup_handlers"):
        context._cleanup_handlers.append(callback)
    else:  # pragma: no cover - fallback safety
        context._cleanup_handlers = [callback]


def _prepare_base_module(context: Any) -> None:
    """Stub external dependencies and load the base agent module."""

    recorded_modules: dict[str, ModuleType | None] = {}

    def set_module(name: str, module: ModuleType) -> ModuleType:
        if name not in recorded_modules:
            recorded_modules[name] = sys.modules.get(name)
        sys.modules[name] = module
        return module

    def ensure_package(name: str) -> ModuleType:
        if name in sys.modules:
            if name not in recorded_modules:
                recorded_modules[name] = sys.modules[name]
            module = sys.modules[name]
        else:
            module = ModuleType(name)
            module.__path__ = []  # type: ignore[attr-defined]
            recorded_modules[name] = None
            sys.modules[name] = module
        return module

    # Stub langchain_core.language_models
    langchain_core = ensure_package("langchain_core")
    language_models = ModuleType("langchain_core.language_models")

    class StubBaseLanguageModel:
        """Placeholder for BaseLanguageModel."""

    language_models.BaseLanguageModel = StubBaseLanguageModel  # type: ignore[attr-defined]
    set_module("langchain_core.language_models", language_models)
    langchain_core.language_models = language_models  # type: ignore[attr-defined]

    # Stub langchain_community.llms
    langchain_community = ensure_package("langchain_community")
    llms_module = ModuleType("langchain_community.llms")

    class StubFakeListLLM:
        """Placeholder for FakeListLLM that tracks instantiations."""

        instances: list[StubFakeListLLM] = []

        def __init__(self, responses: list[str], sleep: float) -> None:
            self.responses = responses
            self.sleep = sleep
            StubFakeListLLM.instances.append(self)

    llms_module.FakeListLLM = StubFakeListLLM  # type: ignore[attr-defined]
    set_module("langchain_community.llms", llms_module)
    langchain_community.llms = llms_module  # type: ignore[attr-defined]

    # Stub langgraph modules
    langgraph = ensure_package("langgraph")
    graph_module = ModuleType("langgraph.graph")

    class StubStateGraph:
        """Placeholder for StateGraph."""

        def __init__(self, state_schema: Any) -> None:
            self.state_schema = state_schema

    graph_module.StateGraph = StubStateGraph  # type: ignore[attr-defined]
    set_module("langgraph.graph", graph_module)
    langgraph.graph = graph_module  # type: ignore[attr-defined]

    checkpoint_pkg = ensure_package("langgraph.checkpoint")
    memory_module = ModuleType("langgraph.checkpoint.memory")

    class StubMemorySaver:
        """Placeholder for MemorySaver that tracks instantiations."""

        instances: list[StubMemorySaver] = []

        def __init__(self) -> None:
            StubMemorySaver.instances.append(self)

    memory_module.MemorySaver = StubMemorySaver  # type: ignore[attr-defined]
    set_module("langgraph.checkpoint.memory", memory_module)
    checkpoint_pkg.memory = memory_module  # type: ignore[attr-defined]

    context.stub_classes = {
        "FakeListLLM": StubFakeListLLM,
        "MemorySaver": StubMemorySaver,
    }

    # Create placeholder package hierarchy for module loading
    package_names = [
        "behave_support",
        "behave_support.cleveragents",
        "behave_support.cleveragents.agents",
    ]
    for package_name in package_names:
        if package_name in sys.modules:
            if package_name not in recorded_modules:
                recorded_modules[package_name] = sys.modules[package_name]
        else:
            pkg = ModuleType(package_name)
            pkg.__path__ = []  # type: ignore[attr-defined]
            set_module(package_name, pkg)

    module_name = "behave_support.cleveragents.agents.base"
    recorded_modules[module_name] = sys.modules.get(module_name)
    spec = importlib.util.spec_from_file_location(module_name, BASE_MODULE_PATH)
    if spec is None or spec.loader is None:  # pragma: no cover - defensive
        raise RuntimeError("Unable to load cleveragents.agents.base module")
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)

    sys.modules["behave_support.cleveragents.agents"].base = module  # type: ignore[attr-defined]

    context.base_module = module
    context.BaseAgent = module.BaseAgent
    context.BaseStateGraph = module.BaseStateGraph

    def cleanup() -> None:
        StubFakeListLLM.instances.clear()
        StubMemorySaver.instances.clear()
        for name, original in recorded_modules.items():
            if original is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = original

    _register_cleanup(context, cleanup)


def _build_input(message: str) -> dict[str, Any]:
    return {
        "messages": [{"role": "user", "content": message}],
        "context": {},
    }


@given("the cleveragents base agent module is loaded with stub dependencies")
def step_load_base_module(context: Any) -> None:
    _prepare_base_module(context)


@given("I instantiate a base agent test double with defaults")
def step_instantiate_test_agent(context: Any) -> None:
    FakeListLLM = context.stub_classes["FakeListLLM"]
    MemorySaver = context.stub_classes["MemorySaver"]
    FakeListLLM.instances.clear()
    MemorySaver.instances.clear()
    context.latest_graph = None
    context.recording_app = None

    BaseAgent = context.BaseAgent

    class TestAgent(BaseAgent):
        def _build_graph(self_inner) -> TestGraph:  # type: ignore[override]
            graph = TestGraph(context)
            context.latest_graph = graph
            return graph

    context.agent = TestAgent()
    context.compiled_app = context.recording_app
    context.invocation_result = None
    context.async_result = None
    context.stream_output = None


@then("the fake list LLM should be constructed once")
def step_verify_fake_llm(context: Any) -> None:
    FakeListLLM = context.stub_classes["FakeListLLM"]
    assert len(FakeListLLM.instances) == 1
    fake_instance = FakeListLLM.instances[0]
    assert fake_instance.responses == ["Mock response"]
    assert fake_instance.sleep == 0.1


@then("the memory saver should be constructed once")
def step_verify_memory_saver(context: Any) -> None:
    MemorySaver = context.stub_classes["MemorySaver"]
    assert len(MemorySaver.instances) == 1


@then("the compiled graph should receive the created memory saver")
def step_verify_compile_checkpointer(context: Any) -> None:
    MemorySaver = context.stub_classes["MemorySaver"]
    assert context.latest_graph is not None
    assert context.latest_graph.compile_checkpointers
    assert context.latest_graph.compile_checkpointers[0] is MemorySaver.instances[0]


@then("the agent should expose a compiled application")
def step_verify_compiled_app(context: Any) -> None:
    assert context.agent.app is context.compiled_app
    assert isinstance(context.compiled_app, RecordingApp)


@when("I call the base agent _build_graph implementation")
def step_call_base_build_graph(context: Any) -> None:
    base_class = context.BaseAgent
    bare_instance = object.__new__(base_class)
    try:
        context.build_graph_error = None
        base_class._build_graph(bare_instance)
    except NotImplementedError as exc:  # pragma: no cover - expected path
        context.build_graph_error = exc


@then("a NotImplementedError should mention subclasses must implement the graph")
def step_verify_not_implemented(context: Any) -> None:
    assert context.build_graph_error is not None
    assert "must implement" in str(context.build_graph_error)


@when('I invoke the test agent with message "{message}" and no config')
def step_invoke_without_config(context: Any, message: str) -> None:
    input_data = _build_input(message)
    context.invocation_result = context.agent.invoke(input_data)


@when(
    'I invoke the test agent with message "{message}" and config thread_id "{thread_id}"'
)
def step_invoke_with_config(context: Any, message: str, thread_id: str) -> None:
    input_data = _build_input(message)
    config = {"configurable": {"thread_id": thread_id}}
    context.invocation_result = context.agent.invoke(input_data, config)


@then('the underlying invoke config should use thread_id "{thread_id}"')
def step_verify_invoke_config(context: Any, thread_id: str) -> None:
    assert context.recording_app.invoke_calls, "No invocation recorded"
    config = context.recording_app.invoke_calls[-1]["config"]
    assert config["configurable"]["thread_id"] == thread_id


@then('the invoke result should include message "{message}"')
def step_verify_invoke_result_message(context: Any, message: str) -> None:
    assert context.invocation_result is not None
    messages = context.invocation_result["input"]["messages"]
    assert messages[0]["content"] == message


@when('I asynchronously invoke the test agent with message "{message}" and no config')
def step_async_invoke_default(context: Any, message: str) -> None:
    input_data = _build_input(message)
    context.async_result = asyncio.run(context.agent.ainvoke(input_data))


@when(
    'I asynchronously invoke the test agent with message "{message}" '
    'and config thread_id "{thread_id}"'
)
def step_async_invoke_custom(context: Any, message: str, thread_id: str) -> None:
    input_data = _build_input(message)
    config = {"configurable": {"thread_id": thread_id}}
    context.async_result = asyncio.run(context.agent.ainvoke(input_data, config))


@then('the async invocation config should use thread_id "{thread_id}"')
def step_verify_async_config(context: Any, thread_id: str) -> None:
    assert context.recording_app.ainvoke_calls, "No async invocation recorded"
    config = context.recording_app.ainvoke_calls[-1]["config"]
    assert config["configurable"]["thread_id"] == thread_id


@then('the async result should include message "{message}"')
def step_verify_async_result(context: Any, message: str) -> None:
    assert context.async_result is not None
    messages = context.async_result["input"]["messages"]
    assert messages[0]["content"] == message


@when('I stream the test agent with message "{message}" and no config')
def step_stream_default(context: Any, message: str) -> None:
    input_data = _build_input(message)
    context.stream_output = list(context.agent.stream(input_data))


@when(
    'I stream the test agent with message "{message}" and config thread_id "{thread_id}"'
)
def step_stream_custom(context: Any, message: str, thread_id: str) -> None:
    input_data = _build_input(message)
    config = {"configurable": {"thread_id": thread_id}}
    context.stream_output = list(context.agent.stream(input_data, config))


@then('the stream events should include thread_id "{thread_id}"')
def step_verify_stream_events(context: Any, thread_id: str) -> None:
    assert context.recording_app.stream_calls, "No stream invocation recorded"
    config = context.recording_app.stream_calls[-1]["config"]
    assert config["configurable"]["thread_id"] == thread_id
    assert context.stream_output is not None
    assert any(
        event.get("config", {}).get("configurable", {}).get("thread_id") == thread_id
        for event in context.stream_output
    )