llama-index

1
"""Joint QA Summary graph."""
2

3
from typing import Optional, Sequence
4

5
from llama_index.legacy.indices.list.base import SummaryIndex
6
from llama_index.legacy.indices.vector_store import VectorStoreIndex
7
from llama_index.legacy.ingestion import run_transformations
8
from llama_index.legacy.query_engine.router_query_engine import RouterQueryEngine
9
from llama_index.legacy.schema import Document
10
from llama_index.legacy.service_context import ServiceContext
11
from llama_index.legacy.storage.storage_context import StorageContext
12
from llama_index.legacy.tools.query_engine import QueryEngineTool
13

14
DEFAULT_SUMMARY_TEXT = "Use this index for summarization queries"
15
DEFAULT_QA_TEXT = (
16
    "Use this index for queries that require retrieval of specific "
17
    "context from documents."
18
)
19

20

21
class QASummaryQueryEngineBuilder:
22
    """Joint QA Summary graph builder.
23

24
    Can build a graph that provides a unified query interface
25
    for both QA and summarization tasks.
26

27
    NOTE: this is a beta feature. The API may change in the future.
28

29
    Args:
30
        docstore (BaseDocumentStore): A BaseDocumentStore to use for storing nodes.
31
        service_context (ServiceContext): A ServiceContext to use for
32
            building indices.
33
        summary_text (str): Text to use for the summary index.
34
        qa_text (str): Text to use for the QA index.
35
        node_parser (NodeParser): A NodeParser to use for parsing.
36

37
    """
38

39
    def __init__(
40
        self,
41
        storage_context: Optional[StorageContext] = None,
42
        service_context: Optional[ServiceContext] = None,
43
        summary_text: str = DEFAULT_SUMMARY_TEXT,
44
        qa_text: str = DEFAULT_QA_TEXT,
45
    ) -> None:
46
        """Init params."""
47
        self._storage_context = storage_context or StorageContext.from_defaults()
48
        self._service_context = service_context or ServiceContext.from_defaults()
49
        self._summary_text = summary_text
50
        self._qa_text = qa_text
51

52
    def build_from_documents(
53
        self,
54
        documents: Sequence[Document],
55
    ) -> RouterQueryEngine:
56
        """Build query engine."""
57
        # parse nodes
58
        nodes = run_transformations(
59
            documents, self._service_context.transformations  # type: ignore
60
        )
61

62
        # ingest nodes
63
        self._storage_context.docstore.add_documents(nodes, allow_update=True)
64

65
        # build indices
66
        vector_index = VectorStoreIndex(
67
            nodes,
68
            service_context=self._service_context,
69
            storage_context=self._storage_context,
70
        )
71
        summary_index = SummaryIndex(
72
            nodes,
73
            service_context=self._service_context,
74
            storage_context=self._storage_context,
75
        )
76

77
        vector_query_engine = vector_index.as_query_engine(
78
            service_context=self._service_context
79
        )
80
        list_query_engine = summary_index.as_query_engine(
81
            service_context=self._service_context,
82
            response_mode="tree_summarize",
83
        )
84

85
        # build query engine
86
        return RouterQueryEngine.from_defaults(
87
            query_engine_tools=[
88
                QueryEngineTool.from_defaults(
89
                    vector_query_engine, description=self._qa_text
90
                ),
91
                QueryEngineTool.from_defaults(
92
                    list_query_engine, description=self._summary_text
93
                ),
94
            ],
95
            service_context=self._service_context,
96
            select_multi=False,
97
        )
98