llama-index

1
"""Utils for pretty print."""
2

3
import textwrap
4
from pprint import pprint
5
from typing import Any, Dict
6

7
from llama_index.legacy.core.response.schema import Response
8
from llama_index.legacy.schema import NodeWithScore
9
from llama_index.legacy.utils import truncate_text
10

11

12
def pprint_metadata(metadata: Dict[str, Any]) -> None:
13
    """Display metadata for jupyter notebook."""
14
    pprint(metadata)
15

16

17
def pprint_source_node(
18
    source_node: NodeWithScore, source_length: int = 350, wrap_width: int = 70
19
) -> None:
20
    """Display source node for jupyter notebook."""
21
    source_text_fmt = truncate_text(
22
        source_node.node.get_content().strip(), source_length
23
    )
24
    print(f"Node ID: {source_node.node.node_id}")
25
    print(f"Similarity: {source_node.score}")
26
    print(textwrap.fill(f"Text: {source_text_fmt}\n", width=wrap_width))
27

28

29
def pprint_response(
30
    response: Response,
31
    source_length: int = 350,
32
    wrap_width: int = 70,
33
    show_source: bool = False,
34
) -> None:
35
    """Pretty print response for jupyter notebook."""
36
    if response.response is None:
37
        response_text = "None"
38
    else:
39
        response_text = response.response.strip()
40

41
    response_text = f"Final Response: {response_text}"
42
    print(textwrap.fill(response_text, width=wrap_width))
43

44
    if show_source:
45
        for ind, source_node in enumerate(response.source_nodes):
46
            print("_" * wrap_width)
47
            print(f"Source Node {ind + 1}/{len(response.source_nodes)}")
48
            pprint_source_node(
49
                source_node, source_length=source_length, wrap_width=wrap_width
50
            )
51