llama-index

1
"""FLARE output parsers."""
2

3
from typing import Any, Callable, Optional
4

5
from llama_index.legacy.query_engine.flare.schema import QueryTask
6
from llama_index.legacy.types import BaseOutputParser
7

8

9
def default_parse_is_done_fn(response: str) -> bool:
10
    """Default parse is done function."""
11
    return "done" in response.lower()
12

13

14
def default_format_done_answer(response: str) -> str:
15
    """Default format done answer."""
16
    return response.replace("done", "").strip()
17

18

19
class IsDoneOutputParser(BaseOutputParser):
20
    """Is done output parser."""
21

22
    def __init__(
23
        self,
24
        is_done_fn: Optional[Callable[[str], bool]] = None,
25
        fmt_answer_fn: Optional[Callable[[str], str]] = None,
26
    ) -> None:
27
        """Init params."""
28
        self._is_done_fn = is_done_fn or default_parse_is_done_fn
29
        self._fmt_answer_fn = fmt_answer_fn or default_format_done_answer
30

31
    def parse(self, output: str) -> Any:
32
        """Parse output."""
33
        is_done = default_parse_is_done_fn(output)
34
        if is_done:
35
            return True, self._fmt_answer_fn(output)
36
        else:
37
            return False, output
38

39
    def format(self, output: str) -> str:
40
        """Format a query with structured output formatting instructions."""
41
        raise NotImplementedError
42

43

44
class QueryTaskOutputParser(BaseOutputParser):
45
    """QueryTask output parser.
46

47
    By default, parses output that contains "[Search(query)]" tags.
48

49
    """
50

51
    def parse(self, output: str) -> Any:
52
        """Parse output."""
53
        query_tasks = []
54
        for idx, char in enumerate(output):
55
            if char == "[":
56
                start_idx = idx
57
            elif char == "]":
58
                end_idx = idx
59
                raw_query_str = output[start_idx + 1 : end_idx]
60
                query_str = raw_query_str.split("(")[1].split(")")[0]
61
                query_tasks.append(QueryTask(query_str, start_idx, end_idx))
62
        return query_tasks
63

64
    def format(self, output: str) -> str:
65
        """Format a query with structured output formatting instructions."""
66
        raise NotImplementedError
67