tests/llm/test_base_llm.py

from string import Template

import pytest

from embedchain.llm.base import BaseLlm, BaseLlmConfig


@pytest.fixture
def base_llm():
    config = BaseLlmConfig()
    return BaseLlm(config=config)


def test_is_get_llm_model_answer_not_implemented(base_llm):
    with pytest.raises(NotImplementedError):
        base_llm.get_llm_model_answer()


def test_is_stream_bool():
    with pytest.raises(ValueError):
        config = BaseLlmConfig(stream="test value")
        BaseLlm(config=config)


def test_template_string_gets_converted_to_Template_instance():
    config = BaseLlmConfig(template="test value $query $context")
    llm = BaseLlm(config=config)
    assert isinstance(llm.config.prompt, Template)


def test_is_get_llm_model_answer_implemented():
    class TestLlm(BaseLlm):
        def get_llm_model_answer(self):
            return "Implemented"

    config = BaseLlmConfig()
    llm = TestLlm(config=config)
    assert llm.get_llm_model_answer() == "Implemented"


def test_stream_response(base_llm):
    answer = ["Chunk1", "Chunk2", "Chunk3"]
    result = list(base_llm._stream_response(answer))
    assert result == answer


def test_append_search_and_context(base_llm):
    context = "Context"
    web_search_result = "Web Search Result"
    result = base_llm._append_search_and_context(context, web_search_result)
    expected_result = "Context\nWeb Search Result: Web Search Result"
    assert result == expected_result


def test_access_search_and_get_results(base_llm, mocker):
    base_llm.access_search_and_get_results = mocker.patch.object(
        base_llm, "access_search_and_get_results", return_value="Search Results"
    )
    input_query = "Test query"
    result = base_llm.access_search_and_get_results(input_query)
    assert result == "Search Results"