Coverage for intelligence_toolkit/tests/unit/AI/test_classes.py: 100%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2# Licensed under the MIT license. See LICENSE file in the project. 

3# 

4import pytest 

5 

6from intelligence_toolkit.AI.classes import LLMCallback, VectorData 

7 

8 

9def test_llm_callback_initialization(): 

10 callback = LLMCallback() 

11 assert callback.response == [] 

12 

13 

14def test_llm_callback_on_llm_new_token(): 

15 callback = LLMCallback() 

16 callback.on_llm_new_token("Hello") 

17 assert callback.response == ["Hello"] 

18 

19 callback.on_llm_new_token(" World") 

20 assert callback.response == ["Hello", " World"] 

21 

22 

23def test_llm_callback_multiple_tokens(): 

24 callback = LLMCallback() 

25 tokens = ["The", " quick", " brown", " fox"] 

26 for token in tokens: 

27 callback.on_llm_new_token(token) 

28 assert callback.response == tokens 

29 

30 

31def test_vector_data_structure(): 

32 # Test that VectorData is a properly defined class with annotations 

33 # VectorData uses type annotations but doesn't set default values 

34 assert hasattr(VectorData, '__annotations__') 

35 assert 'hash' in VectorData.__annotations__ 

36 assert 'text' in VectorData.__annotations__ 

37 assert 'vector' in VectorData.__annotations__ 

38 assert 'additional_details' in VectorData.__annotations__ 

39 

40 

41def test_vector_data_assignment(): 

42 vector_data = VectorData() 

43 vector_data.hash = "test_hash_123" 

44 vector_data.text = "Sample text for testing" 

45 vector_data.vector = [0.1, 0.2, 0.3, 0.4] 

46 vector_data.additional_details = {"key": "value"} 

47 

48 assert vector_data.hash == "test_hash_123" 

49 assert vector_data.text == "Sample text for testing" 

50 assert vector_data.vector == [0.1, 0.2, 0.3, 0.4] 

51 assert vector_data.additional_details == {"key": "value"}