llama-index

1
"""Create LlamaIndex agents."""
2

3
from typing import Any, Optional
4

5
from llama_index.legacy.bridge.langchain import (
6
    AgentExecutor,
7
    AgentType,
8
    BaseCallbackManager,
9
    BaseLLM,
10
    initialize_agent,
11
)
12
from llama_index.legacy.langchain_helpers.agents.toolkits import LlamaToolkit
13

14

15
def create_llama_agent(
16
    toolkit: LlamaToolkit,
17
    llm: BaseLLM,
18
    agent: Optional[AgentType] = None,
19
    callback_manager: Optional[BaseCallbackManager] = None,
20
    agent_path: Optional[str] = None,
21
    agent_kwargs: Optional[dict] = None,
22
    **kwargs: Any,
23
) -> AgentExecutor:
24
    """Load an agent executor given a Llama Toolkit and LLM.
25

26
    NOTE: this is a light wrapper around initialize_agent in langchain.
27

28
    Args:
29
        toolkit: LlamaToolkit to use.
30
        llm: Language model to use as the agent.
31
        agent: A string that specified the agent type to use. Valid options are:
32
            `zero-shot-react-description`
33
            `react-docstore`
34
            `self-ask-with-search`
35
            `conversational-react-description`
36
            `chat-zero-shot-react-description`,
37
            `chat-conversational-react-description`,
38
           If None and agent_path is also None, will default to
39
            `zero-shot-react-description`.
40
        callback_manager: CallbackManager to use. Global callback manager is used if
41
            not provided. Defaults to None.
42
        agent_path: Path to serialized agent to use.
43
        agent_kwargs: Additional key word arguments to pass to the underlying agent
44
        **kwargs: Additional key word arguments passed to the agent executor
45

46
    Returns:
47
        An agent executor
48
    """
49
    llama_tools = toolkit.get_tools()
50
    return initialize_agent(
51
        llama_tools,
52
        llm,
53
        agent=agent,
54
        callback_manager=callback_manager,
55
        agent_path=agent_path,
56
        agent_kwargs=agent_kwargs,
57
        **kwargs,
58
    )
59

60

61
def create_llama_chat_agent(
62
    toolkit: LlamaToolkit,
63
    llm: BaseLLM,
64
    callback_manager: Optional[BaseCallbackManager] = None,
65
    agent_kwargs: Optional[dict] = None,
66
    **kwargs: Any,
67
) -> AgentExecutor:
68
    """Load a chat llama agent given a Llama Toolkit and LLM.
69

70
    Args:
71
        toolkit: LlamaToolkit to use.
72
        llm: Language model to use as the agent.
73
        callback_manager: CallbackManager to use. Global callback manager is used if
74
            not provided. Defaults to None.
75
        agent_kwargs: Additional key word arguments to pass to the underlying agent
76
        **kwargs: Additional key word arguments passed to the agent executor
77

78
    Returns:
79
        An agent executor
80
    """
81
    # chat agent
82
    # TODO: explore chat-conversational-react-description
83
    agent_type = AgentType.CONVERSATIONAL_REACT_DESCRIPTION
84
    return create_llama_agent(
85
        toolkit,
86
        llm,
87
        agent=agent_type,
88
        callback_manager=callback_manager,
89
        agent_kwargs=agent_kwargs,
90
        **kwargs,
91
    )
92