Coverage for tests/test_summarize.py: 94%

1import pytest 

2import mall 

3import polars as pl 

4import pyarrow 

5import shutil 

6import os 

7 

8if os._exists("_test_cache"): 

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

10 

11 

12def test_summarize_prompt(): 

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

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

15 x = df.llm.summarize("x") 

16 assert ( 

17 x["summary"][0] 

18 == "You are a helpful summarization engine. Your answer will contain no no capitalization and no explanations. Return no more than 10 words. The answer is the summary of the following text:\n{}" 

19 ) 

20 

21 

22def test_summarize_max(): 

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

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

25 x = df.llm.summarize("x", max_words=5) 

26 assert ( 

27 x["summary"][0] 

28 == "You are a helpful summarization engine. Your answer will contain no no capitalization and no explanations. Return no more than 5 words. The answer is the summary of the following text:\n{}" 

29 )