Multi-step Tool Use (Agents)

1
# define the `web_search` tool.
2
3
def web_search(query: str) -> list[dict]:
4
  # your code for performing a web search goes here
5
  # return [{
6
  #		"url": "https://en.wikipedia.org/wiki/Ontario",
7
  #		"text": "The capital of Ontario is Toronto, ..."
8
  #	}]
9
10
web_search_tool = {
11
    "type": "function",
12
    "function": {
13
        "name": "web_search",
14
        "description": "performs a web search with the specified query",
15
        "parameters": {
16
            "type": "object",
17
            "properties": {
18
                "query": {
19
                    "type": "str",
20
                    "description": "the query to look up"
21
                }
22
            },
23
            "required": ["query"]
24
        }
25
    }
26
}

1
import json
2
import cohere
3
co = cohere.ClientV2(api_key="<YOUR API KEY>")
4
5
# 1 - Add the user message
6
message = "Who is the mayor of the capital of Ontario?"
7
messages = [{"role": "user", "content": message}]
8
9
# 2 - Model generates tool calls, if any
10
model = "command-r-plus-08-2024"
11
res = co.chat(model=model, messages=messages, tools=[web_search_tool])
12
13
# As long as the model sends back tool_calls,
14
# keep invoking tools and sending the results back to the model
15
while res.message.tool_calls:
16
    print("\nTool plan:")
17
    print(
18
        res.message.tool_plan
19
    )  # This will be an observation and a plan with next steps
20
21
    print("\nTool calls:")
22
    for tc in res.message.tool_calls:
23
        print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
24
25
    messages.append(
26
        {
27
            "role": "assistant",
28
            "tool_calls": res.message.tool_calls,
29
            "tool_plan": res.message.tool_plan,
30
        }
31
    )
32
33
    # 3 - Execute tools based on the tool calls generated by the model
34
    print("\nTool results:")
35
    for tc in res.message.tool_calls:
36
        tool_result = web_search(**json.loads(tc.function.arguments))
37
        print(tool_result)
38
        tool_content = [json.dumps(tool_result)]
39
        messages.append(
40
            {"role": "tool", "tool_call_id": tc.id, "content": tool_content}
41
        )
42
43
    # 4 - Model either generates more tool calls or returns a response
44
    res = co.chat(model=model, messages=messages, tools=[web_search_tool])
45
46
print("\nResponse:")
47
print(res.message.content[0].text)
48
49
if res.message.citations:
50
    print("\nCitations:")
51
    for citation in res.message.citations:
52
        print(citation, "\n")

1
# ! pip install tavily-python --q --disable-pip-version-check
2
3
from tavily import TavilyClient
4
5
tavily_client = TavilyClient(api_key="TAVILY_API_KEY")
6
7
# here's a web search engine
8
def web_search(query: str) -> list[dict]:
9
    response = tavily_client.search(query, max_results=3)["results"]
10
    return {"results": response}
11
12
13
# the LLM is equipped with a description of the web search engine
14
web_search_tool = {
15
    "type": "function",
16
    "function": {
17
        "name": "web_search",
18
        "description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
19
        "parameters": {
20
            "type": "object",
21
            "properties": {
22
                "query": {
23
                    "type": "string",
24
                    "description": "Query to search the internet with",
25
                }
26
            },
27
            "required": ["query"],
28
        },
29
    },
30
}

1
# here's a python console, which can be used to access the spreadsheet, but also more generally to code and plot stuff
2
import io, contextlib
3
4
5
def python_interpreter(code: str) -> list[dict]:
6
    output = io.StringIO()
7
    try:
8
        # Redirect stdout to capture print statements
9
        with contextlib.redirect_stdout(output):
10
            exec(code, globals())
11
    except Exception as e:
12
        return {"error": str(e), "executed_code": code}
13
    # Get stdout
14
    return {"console_output": output.getvalue(), "executed_code": code}
15
    
16
# the LLM is equipped with a description of a python console
17
python_interpreter_tool = {
18
    "type": "function",
19
    "function": {
20
        "name": "python_interpreter",
21
        "description": "Executes python code and returns the result. The code runs in a static sandbox without internet access and without interactive mode, so print output or save output to a file.",
22
        "parameters": {
23
            "type": "object",
24
            "properties": {
25
                "code": {
26
                    "type": "string",
27
                    "description": "Python code to execute"
28
                }
29
            },
30
            "required": ["code"]
31
        }
32
    }
33
}
34
35
functions_map = {
36
    "web_search": web_search,
37
    "python_interpreter": python_interpreter,
38
}

1
message = """What's the age and citizenship of the artists who had the top 3 most streamed songs on Spotify in 2023?
2
3
You have access to a dataset with information about Spotify songs from the past 10 years, located at ./spotify_dataset.csv.
4
You also have access to the internet to search for information not available in the dataset.
5
You must use the dataset when you can, and if stuck you can use the internet.
6
Remember to inspect the dataset and get a list of its columnsto understand its structure before trying to query it. Take it step by step.
7
"""

1
model = "command-r-plus-08-2024"
2
tools = [web_search_tool, python_interpreter_tool]
3
4
# Step 1: get user message
5
print(f"USER MESSAGE:\n{message}")
6
print("="*50)
7
8
messages = [{'role': 'user','content': message}]
9
10
# 2 - Model generates tool calls, if any
11
res = co.chat(model=model,
12
        messages=messages,
13
        tools=tools,
14
        temperature=0)
15
16
# Keep invoking tools as long as the model generates tool calls
17
while res.message.tool_calls:
18
    # Tool plan and tool calls
19
    print("\nTOOL PLAN:")
20
    print(res.message.tool_plan)
21
22
    print("\nTOOL CALLS:")
23
    for tc in res.message.tool_calls:
24
        if tc.function.name == "python_interpreter":
25
            print(f"Tool name: {tc.function.name}")
26
            tool_call_prettified = print("\n".join(f"  {line}" for line_num, line in enumerate(json.loads(tc.function.arguments)["code"].splitlines())))
27
            print(tool_call_prettified)
28
        else:
29
            print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
30
31
    messages.append({'role': 'assistant',
32
                    'tool_calls': res.message.tool_calls,
33
                    'tool_plan': res.message.tool_plan})
34
35
    # 3 - Execute tools based on the tool calls generated by the model
36
    print("\nTOOL RESULTS:")
37
    for tc in res.message.tool_calls:
38
        tool_result = functions_map[tc.function.name](**json.loads(tc.function.arguments))
39
        tool_content = [json.dumps(tool_result)]
40
        print(tool_result, "\n")
41
        
42
        messages.append({"role": "tool",
43
                        "tool_call_id": tc.id,
44
                        "content": tool_content}) 
45
46
    # 4 - Model either generates more tool calls or returns a response
47
    res = co.chat(model=model,
48
                messages=messages,
49
                tools=tools,
50
                temperature=0)
51
    
52
messages.append({"role": "assistant", "content": res.message.content[0].text})
53
54
print("\nRESPONSE:")
55
print(res.message.content[0].text)
56
57
if res.message.citations:
58
    print("\nCITATIONS:")
59
    for citation in res.message.citations:
60
        print(f"Start: {citation.start} | End: {citation.end} | Text: '{citation.text}'")
61
        print("Sources:")
62
        if citation.sources:
63
            for source in citation.sources:
64
                print(source.id)
65
        print("-"*50)