mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
31 lines
948 B
Python
31 lines
948 B
Python
|
"""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
|