rag-chatbot-2

1
from unittest.mock import patch
2

3
import pytest
4
from bot.client.ctransformers_client import CtransformersClient
5
from bot.model.model_settings import ModelType, get_model_setting
6

7

8
@pytest.fixture
9
def cpu_config():
10
    config = {
11
        "top_k": 40,
12
        "top_p": 0.95,
13
        "temperature": 0.7,
14
        "repetition_penalty": 1.1,
15
        "last_n_tokens": 64,
16
        "seed": -1,
17
        "batch_size": 8,
18
        "threads": -1,
19
        "max_new_tokens": 1024,
20
        "stop": None,
21
        "stream": False,
22
        "reset": True,
23
        "context_length": 2048,
24
        "gpu_layers": 0,
25
        "mmap": True,
26
        "mlock": False,
27
    }
28
    return config
29

30

31
@pytest.fixture
32
def valid_model_settings():
33
    model_setting = get_model_setting(ModelType.ZEPHYR.value)
34
    return model_setting
35

36

37
@pytest.fixture
38
def invalid_model_settings():
39
    return get_model_setting(ModelType.OPENCHAT.value)
40

41

42
@pytest.fixture
43
def ctransformers_client(mock_model_folder, valid_model_settings, cpu_config):
44
    with patch.object(valid_model_settings, "config", cpu_config):
45
        return CtransformersClient(mock_model_folder, valid_model_settings)
46

47

48
def test_init_raises_value_error_for_invalid_client_type(mock_model_folder, invalid_model_settings):
49
    with pytest.raises(ValueError):
50
        CtransformersClient(mock_model_folder, invalid_model_settings)
51

52

53
def test_encode_prompt(ctransformers_client):
54
    prompt = "Test prompt"
55
    encoded_prompt = ctransformers_client._encode_prompt(prompt)
56
    assert encoded_prompt is not None
57

58

59
def test_generate_answer(ctransformers_client):
60
    prompt = "Tell me a joke"
61
    generated_answer = ctransformers_client.generate_answer(prompt, max_new_tokens=10)
62
    assert generated_answer is not None
63