Performing Tasks Sequentially with Cohere’s RAG

1
! pip install cohere langchain langchain-community pydantic -qq

1
import json
2
import os
3
import cohere
4
5
from tool_def import (
6
    search_developer_docs,
7
    search_developer_docs_tool,
8
    search_internet,
9
    search_internet_tool,
10
    search_code_examples,
11
    search_code_examples_tool,
12
)
13
14
co = cohere.ClientV2(
15
    "COHERE_API_KEY"
16
)  # Get your free API key: https://dashboard.cohere.com/api-keys
17
18
os.environ["TAVILY_API_KEY"] = (
19
    "TAVILY_API_KEY"  # We'll need the Tavily API key to perform internet search. Get your API key: https://app.tavily.com/home
20
)

1
functions_map = {
2
    "search_developer_docs": search_developer_docs,
3
    "search_internet": search_internet,
4
    "search_code_examples": search_code_examples,
5
}

1
tools = [
2
    search_developer_docs_tool,
3
    search_internet_tool,
4
    search_code_examples_tool,
5
]

1
system_message = """## Task and Context
2
You are an assistant who helps developers use Cohere. You are equipped with a number of tools that can provide different types of information. If you can't find the information you need from one tool, you should try other tools if there is a possibility that they could provide the information you need."""

1
model = "command-a-03-2025"
2
3
4
def run_agent(query, messages=None):
5
    if messages is None:
6
        messages = []
7
8
    if "system" not in {m.get("role") for m in messages}:
9
        messages.append({"role": "system", "content": system_message})
10
11
    # Step 1: get user message
12
    print(f"QUESTION:\n{query}")
13
    print("=" * 50)
14
15
    messages.append({"role": "user", "content": query})
16
17
    # Step 2: Generate tool calls (if any)
18
    response = co.chat(
19
        model=model, messages=messages, tools=tools, temperature=0.3
20
    )
21
22
    while response.message.tool_calls:
23
24
        print("TOOL PLAN:")
25
        print(response.message.tool_plan, "\n")
26
        print("TOOL CALLS:")
27
        for tc in response.message.tool_calls:
28
            print(
29
                f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}"
30
            )
31
        print("=" * 50)
32
33
        messages.append(
34
            {
35
                "role": "assistant",
36
                "tool_calls": response.message.tool_calls,
37
                "tool_plan": response.message.tool_plan,
38
            }
39
        )
40
41
        # Step 3: Get tool results
42
        for tc in response.message.tool_calls:
43
            tool_result = functions_map[tc.function.name](
44
                **json.loads(tc.function.arguments)
45
            )
46
            tool_content = []
47
            for data in tool_result:
48
                tool_content.append(
49
                    {
50
                        "type": "document",
51
                        "document": {"data": json.dumps(data)},
52
                    }
53
                )
54
                # Optional: add an "id" field in the "document" object, otherwise IDs are auto-generated
55
            messages.append(
56
                {
57
                    "role": "tool",
58
                    "tool_call_id": tc.id,
59
                    "content": tool_content,
60
                }
61
            )
62
63
        # Step 4: Generate response and citations
64
        response = co.chat(
65
            model=model,
66
            messages=messages,
67
            tools=tools,
68
            temperature=0.3,
69
        )
70
71
    messages.append(
72
        {
73
            "role": "assistant",
74
            "content": response.message.content[0].text,
75
        }
76
    )
77
78
    # Print final response
79
    print("RESPONSE:")
80
    print(response.message.content[0].text)
81
    print("=" * 50)
82
83
    # Print citations (if any)
84
    verbose_source = (
85
        False  # Change to True to display the contents of a source
86
    )
87
    if response.message.citations:
88
        print("CITATIONS:\n")
89
        for citation in response.message.citations:
90
            print(
91
                f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' "
92
            )
93
            print("Sources:")
94
            for idx, source in enumerate(citation.sources):
95
                print(f"{idx+1}. {source.id}")
96
                if verbose_source:
97
                    print(f"{source.tool_output}")
98
            print("\n")
99
100
    return messages

1
messages = run_agent(
2
    "What's the Cohere feature to reorder search results? Do you have any code examples on that?"
3
)

1
messages = run_agent(
2
    "Who are the CEOs of the companies with the top 3 highest market capitalization."
3
)

1
messages = run_agent(
2
    "Who are the authors of the sentence BERT paper?"
3
)