tests/test_memory_decisions.py

import pytest

from ollie_assistants.assistant.memory_decisions import MemoryDecisionService
from ollie_assistants.llm.types import ChatConfig, ChatMessage, ChatResult


class DecisionProvider:
    provider_id = "decision"

    def __init__(self, text: str | None = None) -> None:
        self.text = text

    async def chat(self, messages: list[ChatMessage], config: ChatConfig) -> ChatResult:
        if self.text is not None:
            return ChatResult(text=self.text)
        if "should_search" in messages[0].content:
            return ChatResult(
                text='{"should_search": true, "query": "name", "reason": "asks name"}'
            )
        return ChatResult(
            text='{"should_write": true, "content": "my name is Rohith", '
            '"kind": "name", "reason": "stable user fact"}'
        )


@pytest.mark.asyncio
async def test_memory_decision_service_parses_search_decision() -> None:
    decision = await MemoryDecisionService(DecisionProvider(), "m").decide_search(
        "what is my name?"
    )

    assert decision.should_search
    assert decision.query == "name"


@pytest.mark.asyncio
async def test_memory_decision_service_parses_write_decision() -> None:
    decision = await MemoryDecisionService(DecisionProvider(), "m").decide_write(
        "my name is Rohith",
        "Nice to meet you.",
    )

    assert decision.should_write
    assert decision.kind == "name"


@pytest.mark.asyncio
async def test_memory_decision_service_treats_primitive_json_as_empty() -> None:
    decision = await MemoryDecisionService(DecisionProvider("false"), "m").decide_search(
        "what is my name?"
    )

    assert not decision.should_search
    assert decision.query == "what is my name?"