llama-index

1
from typing import TYPE_CHECKING, Any, Optional
2

3
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
4

5
if TYPE_CHECKING:
6
    from llama_index.legacy.langchain_helpers.agents.tools import (
7
        LlamaIndexTool,
8
    )
9
from llama_index.legacy.tools.types import AsyncBaseTool, ToolMetadata, ToolOutput
10

11
DEFAULT_NAME = "query_engine_tool"
12
DEFAULT_DESCRIPTION = """Useful for running a natural language query
13
against a knowledge base and get back a natural language response.
14
"""
15

16

17
class QueryEngineTool(AsyncBaseTool):
18
    """Query engine tool.
19

20
    A tool making use of a query engine.
21

22
    Args:
23
        query_engine (BaseQueryEngine): A query engine.
24
        metadata (ToolMetadata): The associated metadata of the query engine.
25
    """
26

27
    def __init__(
28
        self,
29
        query_engine: BaseQueryEngine,
30
        metadata: ToolMetadata,
31
        resolve_input_errors: bool = True,
32
    ) -> None:
33
        self._query_engine = query_engine
34
        self._metadata = metadata
35
        self._resolve_input_errors = resolve_input_errors
36

37
    @classmethod
38
    def from_defaults(
39
        cls,
40
        query_engine: BaseQueryEngine,
41
        name: Optional[str] = None,
42
        description: Optional[str] = None,
43
        resolve_input_errors: bool = True,
44
    ) -> "QueryEngineTool":
45
        name = name or DEFAULT_NAME
46
        description = description or DEFAULT_DESCRIPTION
47

48
        metadata = ToolMetadata(name=name, description=description)
49
        return cls(
50
            query_engine=query_engine,
51
            metadata=metadata,
52
            resolve_input_errors=resolve_input_errors,
53
        )
54

55
    @property
56
    def query_engine(self) -> BaseQueryEngine:
57
        return self._query_engine
58

59
    @property
60
    def metadata(self) -> ToolMetadata:
61
        return self._metadata
62

63
    def call(self, *args: Any, **kwargs: Any) -> ToolOutput:
64
        if args is not None and len(args) > 0:
65
            query_str = str(args[0])
66
        elif kwargs is not None and "input" in kwargs:
67
            # NOTE: this assumes our default function schema of `input`
68
            query_str = kwargs["input"]
69
        elif kwargs is not None and self._resolve_input_errors:
70
            query_str = str(kwargs)
71
        else:
72
            raise ValueError(
73
                "Cannot call query engine without specifying `input` parameter."
74
            )
75

76
        response = self._query_engine.query(query_str)
77
        return ToolOutput(
78
            content=str(response),
79
            tool_name=self.metadata.name,
80
            raw_input={"input": query_str},
81
            raw_output=response,
82
        )
83

84
    async def acall(self, *args: Any, **kwargs: Any) -> ToolOutput:
85
        if args is not None and len(args) > 0:
86
            query_str = str(args[0])
87
        elif kwargs is not None and "input" in kwargs:
88
            # NOTE: this assumes our default function schema of `input`
89
            query_str = kwargs["input"]
90
        elif kwargs is not None and self._resolve_input_errors:
91
            query_str = str(kwargs)
92
        else:
93
            raise ValueError("Cannot call query engine without inputs")
94

95
        response = await self._query_engine.aquery(query_str)
96
        return ToolOutput(
97
            content=str(response),
98
            tool_name=self.metadata.name,
99
            raw_input={"input": query_str},
100
            raw_output=response,
101
        )
102

103
    def as_langchain_tool(self) -> "LlamaIndexTool":
104
        from llama_index.legacy.langchain_helpers.agents.tools import (
105
            IndexToolConfig,
106
            LlamaIndexTool,
107
        )
108

109
        tool_config = IndexToolConfig(
110
            query_engine=self.query_engine,
111
            name=self.metadata.name,
112
            description=self.metadata.description,
113
        )
114
        return LlamaIndexTool.from_tool_config(tool_config=tool_config)
115