OpenJarvis/tests/memory/test_context.py

"""Tests for context injection."""

from __future__ import annotations

import uuid
from typing import Any, Dict, List, Optional

from openjarvis.core.events import EventBus, EventType
from openjarvis.core.types import Message, Role
from openjarvis.memory.store import Fact
from openjarvis.tools.storage._stubs import MemoryBackend, RetrievalResult
from openjarvis.tools.storage.context import (
    ContextConfig,
    build_context_message,
    format_context,
    inject_context,
)

# -- Fake backend for testing ------------------------------------------------


class _FakeMemory(MemoryBackend):
    """In-memory backend that returns pre-set results."""

    backend_id = "fake"

    def __init__(
        self,
        results: Optional[List[RetrievalResult]] = None,
    ) -> None:
        self._results = results or []

    def store(
        self,
        content: str,
        *,
        source: str = "",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> str:
        return uuid.uuid4().hex

    def retrieve(
        self,
        query: str,
        *,
        top_k: int = 5,
        **kwargs: Any,
    ) -> List[RetrievalResult]:
        return self._results[:top_k]

    def delete(self, doc_id: str) -> bool:
        return False

    def clear(self) -> None:
        self._results.clear()


# -- Tests -------------------------------------------------------------------


def test_format_context_with_sources():
    results = [
        RetrievalResult(
            content="Python is great",
            score=1.0,
            source="wiki.md",
        ),
        RetrievalResult(
            content="Java is verbose",
            score=0.8,
            source="notes.txt",
        ),
    ]
    text = format_context(results)
    assert "[Source: wiki.md]" in text
    assert "Python is great" in text
    assert "[Source: notes.txt]" in text


def test_format_context_empty():
    assert format_context([]) == ""


def test_build_context_message_role():
    results = [
        RetrievalResult(content="test", score=1.0, source="s.md"),
    ]
    msg = build_context_message(results)
    assert msg.role == Role.SYSTEM
    assert "knowledge base" in msg.content
    assert "test" in msg.content


def test_inject_context_adds_system_message():
    results = [
        RetrievalResult(
            content="relevant info",
            score=0.9,
            source="doc.md",
        ),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="hello")]
    augmented = inject_context("query", messages, backend)
    assert len(augmented) == 2
    assert augmented[0].role == Role.SYSTEM
    assert "relevant info" in augmented[0].content


def test_inject_context_filters_low_score():
    results = [
        RetrievalResult(content="low score", score=0.01),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="hello")]
    cfg = ContextConfig(min_score=0.1)
    augmented = inject_context(
        "query",
        messages,
        backend,
        config=cfg,
    )
    # Low score filtered out — no context added
    assert len(augmented) == 1


def test_inject_context_respects_max_tokens():
    # Each result has ~100 tokens, max is 150 → only 1 should be included
    content = " ".join(f"word{i}" for i in range(100))
    results = [
        RetrievalResult(content=content, score=1.0, source="a.md"),
        RetrievalResult(content=content, score=0.9, source="b.md"),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="test")]
    cfg = ContextConfig(max_context_tokens=150)
    augmented = inject_context(
        "query",
        messages,
        backend,
        config=cfg,
    )
    assert len(augmented) == 2  # system + user
    # Only one source should be cited
    assert augmented[0].content.count("[Source:") == 1


def test_inject_context_disabled():
    results = [
        RetrievalResult(content="data", score=1.0),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="hello")]
    cfg = ContextConfig(enabled=False)
    augmented = inject_context(
        "query",
        messages,
        backend,
        config=cfg,
    )
    assert len(augmented) == 1


def test_inject_context_no_results_returns_original():
    backend = _FakeMemory([])
    messages = [Message(role=Role.USER, content="hello")]
    augmented = inject_context("query", messages, backend)
    assert augmented is messages


def test_inject_context_adds_auto_memory_facts_without_backend():
    messages = [Message(role=Role.USER, content="What is my favorite color?")]
    facts = [Fact(text="The user's favorite color is blue", source="auto")]

    augmented = inject_context("favorite color", messages, None, facts=facts)

    assert len(augmented) == 2
    assert augmented[0].role == Role.SYSTEM
    assert "remembered from prior conversations" in augmented[0].content
    assert "favorite color is blue" in augmented[0].content


def test_inject_context_prioritizes_newest_facts_within_token_budget():
    messages = [Message(role=Role.USER, content="What do you remember?")]
    facts = [
        Fact(text="old fact uses four tokens"),
        Fact(text="new fact uses four tokens"),
    ]

    augmented = inject_context(
        "remember",
        messages,
        None,
        config=ContextConfig(max_context_tokens=5),
        facts=facts,
    )

    assert "new fact uses four tokens" in augmented[0].content
    assert "old fact uses four tokens" not in augmented[0].content


def test_inject_context_merges_with_existing_system_message():
    messages = [
        Message(role=Role.SYSTEM, content="You are OpenJarvis."),
        Message(role=Role.USER, content="What is my favorite color?"),
    ]
    facts = [Fact(text="The user's favorite color is blue")]

    augmented = inject_context("favorite color", messages, None, facts=facts)

    system_messages = [m for m in augmented if m.role == Role.SYSTEM]
    assert len(system_messages) == 1
    assert "You are OpenJarvis." in system_messages[0].content
    assert "favorite color is blue" in system_messages[0].content
    assert messages[0].content == "You are OpenJarvis."


def test_inject_context_collapses_multiple_system_messages():
    messages = [
        Message(role=Role.SYSTEM, content="Identity."),
        Message(role=Role.SYSTEM, content="Persona."),
        Message(role=Role.USER, content="What do you remember?"),
    ]

    augmented = inject_context(
        "remember",
        messages,
        None,
        facts=[Fact(text="User likes jazz")],
    )

    system_messages = [m for m in augmented if m.role == Role.SYSTEM]
    assert len(system_messages) == 1
    assert "Identity." in system_messages[0].content
    assert "Persona." in system_messages[0].content
    assert "User likes jazz" in system_messages[0].content


def test_inject_context_reserves_budget_for_retrieved_documents():
    backend = _FakeMemory(
        [RetrievalResult(content="d1 d2 d3 d4 d5", score=1.0, source="doc")]
    )
    facts = [
        Fact(text="old1 old2 old3 old4 old5"),
        Fact(text="new1 new2 new3 new4 new5"),
    ]

    augmented = inject_context(
        "query",
        [Message(role=Role.USER, content="query")],
        backend,
        config=ContextConfig(max_context_tokens=10),
        facts=facts,
    )

    assert "new1 new2 new3 new4 new5" in augmented[0].content
    assert "d1 d2 d3 d4 d5" in augmented[0].content
    assert "old1 old2 old3 old4 old5" not in augmented[0].content


def test_inject_context_prefers_large_document_that_fits_total_budget():
    backend = _FakeMemory(
        [RetrievalResult(content="d1 d2 d3 d4 d5 d6 d7 d8", score=1.0)]
    )

    augmented = inject_context(
        "query",
        [Message(role=Role.USER, content="query")],
        backend,
        config=ContextConfig(max_context_tokens=10),
        facts=[Fact(text="f1 f2 f3 f4 f5")],
    )

    assert "d1 d2 d3 d4 d5 d6 d7 d8" in augmented[0].content
    assert "f1 f2 f3 f4 f5" not in augmented[0].content


def test_inject_context_publishes_event():
    bus = EventBus(record_history=True)
    results = [
        RetrievalResult(content="info", score=0.9, source="s.md"),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="hello")]

    import openjarvis.tools.storage.context as mod

    original = mod.get_event_bus
    mod.get_event_bus = lambda: bus
    try:
        inject_context("query", messages, backend)
        events = [e for e in bus.history if e.event_type == EventType.MEMORY_RETRIEVE]
        assert len(events) == 1
        assert events[0].data["context_injection"] is True
    finally:
        mod.get_event_bus = original


def test_inject_context_does_not_mutate_original():
    results = [
        RetrievalResult(content="info", score=0.9, source="s.md"),
    ]
    backend = _FakeMemory(results)
    messages = [Message(role=Role.USER, content="hello")]
    original_len = len(messages)
    augmented = inject_context("query", messages, backend)
    assert len(messages) == original_len
    assert len(augmented) == original_len + 1