learningcircuit--local-deep-research/tests/test_llm/test_llm_registry.py

"""Tests for the LLM registry module."""

import pytest
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage
from langchain_core.outputs import ChatGeneration, ChatResult

from local_deep_research.llm import (
    clear_llm_registry,
    get_llm_from_registry,
    is_llm_registered,
    list_registered_llms,
    register_llm,
    unregister_llm,
)


class MockLLM(BaseChatModel):
    """Mock LLM for testing."""

    name: str = "mock"

    def _generate(self, messages, **kwargs):
        """Generate mock response."""
        message = AIMessage(content=f"Mock response from {self.name}")
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @property
    def _llm_type(self):
        return "mock"


@pytest.fixture(autouse=True)
def clear_registry():
    """Clear the registry before and after each test."""
    clear_llm_registry()
    yield
    clear_llm_registry()


def test_register_llm_instance():
    """Test registering an LLM instance."""
    llm = MockLLM(name="test-llm")

    # Register the LLM
    register_llm("test_model", llm)

    # Check it's registered
    assert is_llm_registered("test_model")
    assert "test_model" in list_registered_llms()

    # Retrieve it
    retrieved = get_llm_from_registry("test_model")
    assert retrieved is llm


def test_register_llm_factory():
    """Test registering an LLM factory function."""

    def create_llm(**kwargs):
        return MockLLM(name="factory-llm", **kwargs)

    # Register the factory
    register_llm("factory_model", create_llm)

    # Check it's registered
    assert is_llm_registered("factory_model")

    # Retrieve the factory
    factory = get_llm_from_registry("factory_model")
    assert callable(factory)

    # Create an LLM from the factory
    llm = factory()
    assert isinstance(llm, MockLLM)
    assert llm.name == "factory-llm"


def test_unregister_llm():
    """Test unregistering an LLM."""
    llm = MockLLM()
    register_llm("temp_model", llm)

    # Verify it's registered
    assert is_llm_registered("temp_model")

    # Unregister it
    unregister_llm("temp_model")

    # Verify it's gone
    assert not is_llm_registered("temp_model")
    assert "temp_model" not in list_registered_llms()
    assert get_llm_from_registry("temp_model") is None


def test_multiple_llms():
    """Test registering multiple LLMs."""
    llm1 = MockLLM(name="llm1")
    llm2 = MockLLM(name="llm2")

    register_llm("model1", llm1)
    register_llm("model2", llm2)

    # Check both are registered
    registered = list_registered_llms()
    assert len(registered) == 2
    assert "model1" in registered
    assert "model2" in registered

    # Retrieve them
    assert get_llm_from_registry("model1") is llm1
    assert get_llm_from_registry("model2") is llm2


def test_overwrite_existing():
    """Test overwriting an existing LLM."""
    llm1 = MockLLM(name="original")
    llm2 = MockLLM(name="replacement")

    # Register first LLM
    register_llm("model", llm1)
    assert get_llm_from_registry("model") is llm1

    # Overwrite with second LLM
    register_llm("model", llm2)
    assert get_llm_from_registry("model") is llm2


def test_clear_registry():
    """Test clearing all registered LLMs."""
    # Register multiple LLMs
    register_llm("model1", MockLLM())
    register_llm("model2", MockLLM())
    register_llm("model3", MockLLM())

    assert len(list_registered_llms()) == 3

    # Clear the registry
    clear_llm_registry()

    # Verify all are gone
    assert len(list_registered_llms()) == 0
    assert not is_llm_registered("model1")
    assert not is_llm_registered("model2")
    assert not is_llm_registered("model3")


def test_thread_safety():
    """Test that registry operations are thread-safe."""
    import threading
    import time

    results = []
    errors = []

    def register_many():
        """Register many LLMs in a thread."""
        try:
            for i in range(100):
                register_llm(f"thread_model_{i}", MockLLM(name=f"thread-{i}"))
            results.append("register_complete")
        except Exception as e:
            errors.append(e)

    def read_many():
        """Read from registry in a thread."""
        try:
            for _ in range(100):
                _ = list_registered_llms()
                time.sleep(0.001)  # Small delay to increase contention
            results.append("read_complete")
        except Exception as e:
            errors.append(e)

    # Start multiple threads
    threads = []
    for _ in range(3):
        threads.append(threading.Thread(target=register_many))
        threads.append(threading.Thread(target=read_many))

    for t in threads:
        t.start()

    for t in threads:
        t.join()

    # Check no errors occurred
    assert len(errors) == 0
    assert len(results) == 6  # 3 register + 3 read threads


def test_get_nonexistent():
    """Test getting a non-existent LLM."""
    assert get_llm_from_registry("does_not_exist") is None
    assert not is_llm_registered("does_not_exist")


def test_empty_registry():
    """Test operations on empty registry."""
    assert list_registered_llms() == []
    assert not is_llm_registered("any_model")
    assert get_llm_from_registry("any_model") is None