Coverage for tests/test_gen.py: 76%

1from typing import List 

2 

3import pytest 

4 

5import appl 

6from appl import ( 

7 AIRole, 

8 CompletionResponse, 

9 PromptRecords, 

10 SystemMessage, 

11 gen, 

12 ppl, 

13 records, 

14) 

15 

16 

17def test_gen(): 

18 appl.init() 

19 

20 @ppl 

21 def func(): 

22 "Hello" 

23 return gen("_dummy", mock_response="World") 

24 

25 assert str(func()) == "World" 

26 

27 @ppl 

28 def dummy(): 

29 "Hello" 

30 return gen("_dummy") 

31 

32 assert "dummy" in str(dummy()) 

33 

34 

35def test_chat_gen(): 

36 appl.init() 

37 name = "Test" 

38 p: List[PromptRecords] = [] 

39 

40 @ppl(auto_prime=True) 

41 def chat_gen(): 

42 SystemMessage(f"Your name is {name}") 

43 reply = None 

44 while True: 

45 yield reply 

46 with AIRole(): 

47 (reply := gen("_dummy", mock_response="World")) 

48 p.append(records()) 

49 

50 chat = chat_gen() 

51 assert str(chat.send("Hello")) == "World" 

52 print(p[0].as_convo().as_list()) 

53 assert p[0].as_convo().as_list() == [ 

54 {"role": "system", "content": "Your name is Test"}, 

55 {"role": "user", "content": "Hello"}, 

56 {"role": "assistant", "content": "World"}, 

57 ] 

58 assert str(chat.send("Hello again")) == "World" 

59 assert p[1].as_convo().as_list() == [ 

60 {"role": "system", "content": "Your name is Test"}, 

61 {"role": "user", "content": "Hello"}, 

62 {"role": "assistant", "content": "World"}, 

63 {"role": "user", "content": "Hello again"}, 

64 {"role": "assistant", "content": "World"}, 

65 ] 

66 

67 

68def test_gen_error(): 

69 with pytest.raises(ValueError) as excinfo: 

70 gen(mock_response="World") 

71 

72 error_msg = str(excinfo.value) 

73 assert "PromptContext" in error_msg or "appl.init" in error_msg