File size: 948 Bytes
58d33f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
"""Test LLM callbacks."""
from langchain.callbacks.base import CallbackManager
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
from tests.unit_tests.llms.fake_llm import FakeLLM


def test_llm_with_callbacks() -> None:
    """Test LLM callbacks."""
    handler = FakeCallbackHandler()
    llm = FakeLLM(callback_manager=CallbackManager(handlers=[handler]), verbose=True)
    output = llm("foo")
    assert output == "foo"
    assert handler.starts == 1
    assert handler.ends == 1
    assert handler.errors == 0


def test_llm_with_callbacks_not_verbose() -> None:
    """Test LLM callbacks but not verbose."""
    import langchain

    langchain.verbose = False

    handler = FakeCallbackHandler()
    llm = FakeLLM(callback_manager=CallbackManager(handlers=[handler]))
    output = llm("foo")
    assert output == "foo"
    assert handler.starts == 0
    assert handler.ends == 0
    assert handler.errors == 0