Coverage for tests/test_extract.py: 96%

1 import pytest

2 import mall

3 import polars as pl

4 import pyarrow

5

6 import shutil

7 import os

8 if os._exists("_test_cache"):

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

10

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

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

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