Tool Use & Agents

1
# %pip install cohere
2
import cohere
3
4
co = cohere.ClientV2(
5
    api_key="AZURE_API_KEY_CHAT",
6
    base_url="AZURE_ENDPOINT_CHAT",  # example: "https://cohere-command-r-plus-08-2024-xyz.eastus.models.ai.azure.com/"
7
)

1
def search_faqs(query):
2
    faqs = [
3
        {
4
            "text": "Submitting Travel Expenses:\nSubmit your expenses through our user-friendly finance tool."
5
        },
6
        {
7
            "text": "Side Projects Policy:\nWe encourage you to explore your passions! Just ensure there's no conflict of interest with our business."
8
        },
9
        {
10
            "text": "Wellness Benefits:\nTo promote a healthy lifestyle, we provide gym memberships, on-site yoga classes, and health insurance."
11
        },
12
    ]
13
    return faqs
14
15
16
def search_emails(query):
17
    emails = [
18
        {
19
            "from": "hr@co1t.com",
20
            "to": "david@co1t.com",
21
            "date": "2024-06-24",
22
            "subject": "A Warm Welcome to Co1t, David!",
23
            "text": "We are delighted to have you on board. Please find attached your first week's agenda.",
24
        },
25
        {
26
            "from": "it@co1t.com",
27
            "to": "david@co1t.com",
28
            "date": "2024-06-24",
29
            "subject": "Instructions for IT Setup",
30
            "text": "Welcome, David! To get you started, please follow the attached guide to set up your work accounts.",
31
        },
32
        {
33
            "from": "john@co1t.com",
34
            "to": "david@co1t.com",
35
            "date": "2024-06-24",
36
            "subject": "First Week Check-In",
37
            "text": "Hi David, let's chat briefly tomorrow to discuss your first week. Also, come join us for lunch this Thursday at noon to meet everyone!",
38
        },
39
    ]
40
    return emails
41
42
43
def create_calendar_event(date: str, time: str, duration: int):
44
    # You can implement any logic here
45
    return {
46
        "is_success": True,
47
        "message": f"Created a {duration} hour long event at {time} on {date}",
48
    }
49
50
51
functions_map = {
52
    "search_faqs": search_faqs,
53
    "search_emails": search_emails,
54
    "create_calendar_event": create_calendar_event,
55
}

1
tools = [
2
    {
3
        "type": "function",
4
        "function": {
5
            "name": "search_faqs",
6
            "description": "Given a user query, searches a company's frequently asked questions (FAQs) list and returns the most relevant matches to the query.",
7
            "parameters": {
8
                "type": "object",
9
                "properties": {
10
                    "query": {
11
                        "type": "string",
12
                        "description": "The query from the user",
13
                    }
14
                },
15
                "required": ["query"],
16
            },
17
        },
18
    },
19
    {
20
        "type": "function",
21
        "function": {
22
            "name": "search_emails",
23
            "description": "Given a user query, searches a person's emails and returns the most relevant matches to the query.",
24
            "parameters": {
25
                "type": "object",
26
                "properties": {
27
                    "query": {
28
                        "type": "string",
29
                        "description": "The query from the user",
30
                    }
31
                },
32
                "required": ["query"],
33
            },
34
        },
35
    },
36
    {
37
        "type": "function",
38
        "function": {
39
            "name": "create_calendar_event",
40
            "description": "Creates a new calendar event of the specified duration at the specified time and date. A new event cannot be created on the same time as an existing event.",
41
            "parameters": {
42
                "type": "object",
43
                "properties": {
44
                    "date": {
45
                        "type": "string",
46
                        "description": "the date on which the event starts, formatted as mm/dd/yy",
47
                    },
48
                    "time": {
49
                        "type": "string",
50
                        "description": "the time of the event, formatted using 24h military time formatting",
51
                    },
52
                    "duration": {
53
                        "type": "number",
54
                        "description": "the number of hours the event lasts for",
55
                    },
56
                },
57
                "required": ["date", "time", "duration"],
58
            },
59
        },
60
    },
61
]

1
import json
2
3
system_message = """## Task and Context
4
You are an assistant who assists new employees of Co1t with their first week. You respond to their questions and assist them with their needs. Today is Monday, June 24, 2024"""
5
6
7
def run_assistant(query, messages=None):
8
    if messages is None:
9
        messages = []
10
11
    if "system" not in {m.get("role") for m in messages}:
12
        messages.append({"role": "system", "content": system_message})
13
14
    # Step 1: get user message
15
    print(f"Question:\n{query}")
16
    print("=" * 50)
17
18
    messages.append({"role": "user", "content": query})
19
20
    # Step 2: Generate tool calls (if any)
21
    response = co.chat(
22
        model="model",  # Pass a dummy string
23
        messages=messages,
24
        tools=tools,
25
    )
26
27
    while response.message.tool_calls:
28
29
        print("Tool plan:")
30
        print(response.message.tool_plan, "\n")
31
        print("Tool calls:")
32
        for tc in response.message.tool_calls:
33
            print(
34
                f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}"
35
            )
36
        print("=" * 50)
37
38
        messages.append(
39
            {
40
                "role": "assistant",
41
                "tool_calls": response.message.tool_calls,
42
                "tool_plan": response.message.tool_plan,
43
            }
44
        )
45
46
        # Step 3: Get tool results
47
        for idx, tc in enumerate(response.message.tool_calls):
48
            tool_result = functions_map[tc.function.name](
49
                **json.loads(tc.function.arguments)
50
            )
51
            tool_content = []
52
            for data in tool_result:
53
                tool_content.append(
54
                    {
55
                        "type": "document",
56
                        "document": {"data": json.dumps(data)},
57
                    }
58
                )
59
                # Optional: add an "id" field in the "document" object, otherwise IDs are auto-generated
60
            messages.append(
61
                {
62
                    "role": "tool",
63
                    "tool_call_id": tc.id,
64
                    "content": tool_content,
65
                }
66
            )
67
68
        # Step 4: Generate response and citations
69
        response = co.chat(
70
            model="model",  # Pass a dummy string
71
            messages=messages,
72
            tools=tools,
73
        )
74
75
    messages.append(
76
        {
77
            "role": "assistant",
78
            "content": response.message.content[0].text,
79
        }
80
    )
81
82
    # Print final response
83
    print("Response:")
84
    print(response.message.content[0].text)
85
    print("=" * 50)
86
87
    # Print citations (if any)
88
    if response.message.citations:
89
        print("\nCITATIONS:")
90
        for citation in response.message.citations:
91
            print(citation, "\n")
92
93
    return messages

1
messages = run_assistant(
2
    "Any doc on how do I submit travel expenses? Also, any emails about setting up IT access?"
3
)