Coverage for tests/test_summarize.py: 94%

1 import pytest

2 import mall

3 import polars as pl

4 import pyarrow

5 import shutil

6 import os

7

8 if os._exists("_test_cache"):

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

10

11

12 def 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

22 def 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 )