Coverage for tests/test_extract.py: 96%

1import pytest 

2import mall 

3import polars as pl 

4import pyarrow 

5 

6import shutil 

7import os 

8if os._exists("_test_cache"): 

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

10 

11def test_extract_list(): 

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

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

14 x = df.llm.extract("x", ["a", "b"]) 

15 assert ( 

16 x["extract"][0] 

17 == "You are a helpful text extraction engine. Extract the a, b being referred to on the text. I expect 2 items exactly. No capitalization. No explanations. Return the response exclusively in a pipe separated list, and no headers. The answer is based on the following text:\n{}" 

18 ) 

19 

20 

21def test_extract_dict(): 

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

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

24 x = df.llm.extract("x", dict(a="one", b="two")) 

25 assert ( 

26 x["extract"][0] 

27 == "You are a helpful text extraction engine. Extract the one, two being referred to on the text. I expect 2 items exactly. No capitalization. No explanations. Return the response exclusively in a pipe separated list, and no headers. The answer is based on the following text:\n{}" 

28 ) 

29 

30 

31def test_extract_one(): 

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

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

34 x = df.llm.extract("x", labels="a") 

35 assert ( 

36 x["extract"][0] 

37 == "You are a helpful text extraction engine. Extract the a being referred to on the text. I expect 1 item exactly. No capitalization. No explanations. The answer is based on the following text:\n{}" 

38 )