cleveragents-core/robot/helper_session_tell_llm.py

"""Helper script for session_tell_llm.robot integration tests.

Each subcommand is a self-contained check that prints a sentinel on success.
Tests verify SessionWorkflow.tell() / tell_stream() behaviour with a stub LLM.

No real API keys or network access required — all LLM calls are intercepted
by a monkey-patched ``_resolve_llm`` that returns a minimal stub.
"""

# ruff: noqa: E402, I001
from __future__ import annotations

import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Any

# Ensure local source tree is importable
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from cleveragents.application.services.session_service import (
    PersistentSessionService,
)
from cleveragents.application.services.session_workflow import (
    SessionWorkflow,
)
from cleveragents.domain.models.core.session import (
    MessageRole,
    SessionActorNotConfiguredError,
)
from cleveragents.infrastructure.database.models import Base
from cleveragents.infrastructure.database.repositories import (
    SessionMessageRepository,
    SessionRepository,
)


# ---------------------------------------------------------------------------
# Stub LLM — returns deterministic responses without LLM API calls
# ---------------------------------------------------------------------------

_STUB_TEXT = "This is a real response from the stub LLM actor."


class _StubLLM:
    """Minimal LLM stub that behaves like a LangChain chat model."""

    def invoke(self, messages: Any, **kwargs: Any) -> Any:
        class _Resp:
            content = _STUB_TEXT

            def __init__(self) -> None:
                self.tool_calls: list[Any] = []
                self.response_metadata: dict[str, Any] = {
                    "usage": {"input_tokens": 10, "output_tokens": 20}
                }

        return _Resp()

    def stream(self, messages: Any, **kwargs: Any) -> Iterator[Any]:
        for word in _STUB_TEXT.split():

            class _Chunk:
                def __init__(self, t: str) -> None:
                    self.content = t + " "

            yield _Chunk(word)


# ---------------------------------------------------------------------------
# Fixture helpers
# ---------------------------------------------------------------------------


def _make_db_session_factory() -> Any:
    """Create an in-memory SQLite engine and return a SQLAlchemy session factory."""
    engine = create_engine("sqlite:///:memory:", echo=False)
    Base.metadata.create_all(engine)
    return sessionmaker(bind=engine, expire_on_commit=False)


def _make_workflow(
    actor_name: str | None = "openai/gpt-4",
    llm_factory: Any = None,
) -> tuple[SessionWorkflow, Any]:
    """Build a real SessionWorkflow backed by an in-memory SQLite database."""
    db_factory = _make_db_session_factory()
    db_session = db_factory()

    def get_db() -> Any:
        return db_session

    session_repo = SessionRepository(get_db)
    message_repo = SessionMessageRepository(get_db)
    svc = PersistentSessionService(session_repo=session_repo, message_repo=message_repo)

    # Create a real session in the DB
    session = svc.create(actor_name=actor_name)

    factory = llm_factory if llm_factory is not None else lambda _actor: _StubLLM()
    wf = SessionWorkflow(session_service=svc, llm_factory=factory)
    # No monkey-patching needed — llm_factory provides clean DI (addresses M8)

    return wf, session


# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------


def tell_persists() -> None:
    """Verify tell() returns real LLM response and persists messages."""
    wf, session = _make_workflow(actor_name="openai/gpt-4")
    result = wf.tell(
        session_id=session.session_id,
        prompt="What can you help me with?",
    )
    assert result.assistant_message == _STUB_TEXT, (
        f"Expected stub response, got: {result.assistant_message!r}"
    )
    assert result.input_tokens > 0, "input_tokens should be > 0"
    assert result.output_tokens > 0, "output_tokens should be > 0"
    # Verify messages were persisted via the service
    messages = wf._session_service.get_messages(session.session_id)
    roles = [m.role for m in messages]
    assert MessageRole.USER in roles, "User message not persisted"
    assert MessageRole.ASSISTANT in roles, "Assistant message not persisted"
    print("session-tell-persists-ok")


def tell_no_actor() -> None:
    """Verify tell() raises SessionActorNotConfiguredError when no actor set."""
    wf, session = _make_workflow(actor_name=None)
    try:
        wf.tell(session_id=session.session_id, prompt="Hello")
        print(
            "FAIL: expected SessionActorNotConfiguredError not raised",
            file=sys.stderr,
        )
        sys.exit(1)
    except SessionActorNotConfiguredError as exc:
        assert "actor" in str(exc).lower(), (
            f"Error message should mention 'actor', got: {exc}"
        )
        print("session-tell-no-actor-ok")


def tell_stream() -> None:
    """Verify tell_stream() yields tokens from the stub LLM."""
    wf, session = _make_workflow(actor_name="openai/gpt-4")
    tokens: list[str] = []
    for token in wf.tell_stream(
        session_id=session.session_id,
        prompt="Hello",
    ):
        tokens.append(token)
    assert len(tokens) > 0, "Expected tokens to be yielded, got none"
    full_response = "".join(tokens).strip()
    assert len(full_response) > 0, f"Expected non-empty stream, got: {full_response!r}"
    print("session-tell-stream-ok")


def tell_actor_override() -> None:
    """Verify tell() uses actor_override instead of session actor."""
    captured_actor: list[str] = []

    def _tracking_resolve(actor_name: str) -> Any:
        captured_actor.append(actor_name)
        return _StubLLM()

    wf, session = _make_workflow(
        actor_name="anthropic/claude-3-haiku",
        llm_factory=_tracking_resolve,
    )

    result = wf.tell(
        session_id=session.session_id,
        prompt="Hello",
        actor_override="openai/gpt-4",
    )
    assert captured_actor, "Actor resolver was never called"
    assert captured_actor[0] == "openai/gpt-4", (
        f"Expected 'openai/gpt-4' actor, got {captured_actor[0]!r}"
    )
    assert result.assistant_message is not None
    print("session-tell-actor-override-ok")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

COMMANDS = {
    "tell-persists": tell_persists,
    "tell-no-actor": tell_no_actor,
    "tell-stream": tell_stream,
    "tell-actor-override": tell_actor_override,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        print(
            f"Usage: {sys.argv[0]} <command>\nCommands: {', '.join(COMMANDS)}",
            file=sys.stderr,
        )
        sys.exit(1)
    COMMANDS[sys.argv[1]]()