Coverage for tests/test_sentiment.py: 97%

1import pytest 

2import mall 

3import polars as pl 

4import pyarrow 

5 

6import shutil 

7import os 

8 

9if os._exists("_test_cache"): 

10 shutil.rmtree("_test_cache", ignore_errors=True) 

11 

12 

13def test_sentiment_simple(): 

14 data = mall.MallData 

15 reviews = data.reviews 

16 reviews.llm.use("test", "echo", _cache="_test_cache") 

17 x = reviews.llm.sentiment("review") 

18 assert ( 

19 x.select("sentiment").to_pandas().to_string() 

20 == " sentiment\n0 None\n1 None\n2 None" 

21 ) 

22 

23 

24def sim_sentiment(): 

25 df = pl.DataFrame(dict(x=["positive", "negative", "neutral", "not-real"])) 

26 df.llm.use("test", "echo", _cache="_test_cache") 

27 return df 

28 

29 

30def test_sentiment_valid(): 

31 x = sim_sentiment() 

32 x = x.llm.sentiment("x") 

33 assert ( 

34 x.select("sentiment").to_pandas().to_string() 

35 == " sentiment\n0 positive\n1 negative\n2 neutral\n3 None" 

36 ) 

37 

38 

39def test_sentiment_valid2(): 

40 x = sim_sentiment() 

41 x = x.llm.sentiment("x", ["positive", "negative"]) 

42 assert ( 

43 x.select("sentiment").to_pandas().to_string() 

44 == " sentiment\n0 positive\n1 negative\n2 None\n3 None" 

45 ) 

46 

47 

48def test_sentiment_prompt(): 

49 df = pl.DataFrame(dict(x="x")) 

50 df.llm.use("test", "content", _cache="_test_cache") 

51 x = df.llm.sentiment("x") 

52 assert ( 

53 x["sentiment"][0] 

54 == "You are a helpful sentiment engine. Return only one of the following answers: positive, negative, neutral . No capitalization. No explanations. The answer is based on the following text:\n{}" 

55 )