llama-index

1
"""LlamaIndex toolkit."""
2

3
from typing import List
4

5
from llama_index.legacy.bridge.langchain import BaseTool, BaseToolkit
6
from llama_index.legacy.bridge.pydantic import Field
7
from llama_index.legacy.langchain_helpers.agents.tools import (
8
    IndexToolConfig,
9
    LlamaIndexTool,
10
)
11

12

13
class LlamaToolkit(BaseToolkit):
14
    """Toolkit for interacting with Llama indices."""
15

16
    index_configs: List[IndexToolConfig] = Field(default_factory=list)
17

18
    class Config:
19
        """Configuration for this pydantic object."""
20

21
        arbitrary_types_allowed = True
22

23
    def get_tools(self) -> List[BaseTool]:
24
        """Get the tools in the toolkit."""
25
        index_tools: List[BaseTool] = [
26
            LlamaIndexTool.from_tool_config(tool_config=tool_config)
27
            for tool_config in self.index_configs
28
        ]
29

30
        return index_tools
31