Basic Tool Use

1
import cohere, json
2
API_KEY = "..." # fill in your Cohere API key here
3
co = cohere.Client(API_KEY)

1
sales_database = {
2
    '2023-09-28': {
3
        'total_sales_amount': 5000,
4
        'total_units_sold': 100,
5
    },
6
    '2023-09-29': {
7
        'total_sales_amount': 10000,
8
        'total_units_sold': 250,
9
    },
10
    '2023-09-30': {
11
        'total_sales_amount': 8000,
12
        'total_units_sold': 200,
13
    }
14
}
15
16
product_catalog = {
17
    'Electronics': [
18
        {'product_id': 'E1001', 'name': 'Smartphone', 'price': 500, 'stock_level': 20},
19
        {'product_id': 'E1002', 'name': 'Laptop', 'price': 1000, 'stock_level': 15},
20
        {'product_id': 'E1003', 'name': 'Tablet', 'price': 300, 'stock_level': 25},
21
    ],
22
    'Clothing': [
23
        {'product_id': 'C1001', 'name': 'T-Shirt', 'price': 20, 'stock_level': 100},
24
        {'product_id': 'C1002', 'name': 'Jeans', 'price': 50, 'stock_level': 80},
25
        {'product_id': 'C1003', 'name': 'Jacket', 'price': 100, 'stock_level': 40},
26
    ]
27
}

1
def query_daily_sales_report(day: str) -> dict:
2
    """
3
    Function to retrieve the sales report for the given day
4
    """
5
    report = sales_database.get(day, {})
6
    if report:
7
        return {
8
            'date': day,
9
            'summary': f"Total Sales Amount: {report['total_sales_amount']}, Total Units Sold: {report['total_units_sold']}"
10
        }
11
    else:
12
        return {'date': day, 'summary': 'No sales data available for this day.'}
13
14
15
def query_product_catalog(category: str) -> dict:
16
    """
17
    Function to retrieve products for the given category
18
    """
19
    products = product_catalog.get(category, [])
20
    return {
21
        'category': category,
22
        'products': products
23
    }
24
25
26
functions_map = {
27
    "query_daily_sales_report": query_daily_sales_report,
28
    "query_product_catalog": query_product_catalog
29
}

1
tools = [
2
    {
3
        "name": "query_daily_sales_report",
4
        "description": "Connects to a database to retrieve overall sales volumes and sales information for a given day.",
5
        "parameter_definitions": {
6
            "day": {
7
                "description": "Retrieves sales data for this day, formatted as YYYY-MM-DD.",
8
                "type": "str",
9
                "required": True
10
            }
11
        }
12
    },
13
    {
14
        "name": "query_product_catalog",
15
        "description": "Connects to a a product catalog with information about all the products being sold, including categories, prices, and stock levels.",
16
        "parameter_definitions": {
17
            "category": {
18
                "description": "Retrieves product information data for all products in this category.",
19
                "type": "str",
20
                "required": True
21
            }
22
        }
23
    }
24
]

1
preamble = """
2
## Task & Context
3
You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.
4
5
## Style Guide
6
Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.
7
"""
8
9
message = "Can you provide a sales summary for 29th September 2023, and also give me some details about the products in the 'Electronics' category, for example their prices and stock levels?"

1
response = co.chat(
2
    message=message,
3
    tools=tools,
4
    preamble=preamble,
5
    model="command-r"
6
)
7
8
9
print("The model recommends doing the following tool calls:")
10
print("\n".join(str(tool_call) for tool_call in response.tool_calls))

1
tool_results = []
2
for tool_call in response.tool_calls:
3
    # here is where you would call the tool recommended by the model, using the parameters recommended by the model
4
    print(f"= running tool {tool_call.name}, with parameters: {tool_call.parameters}")
5
    output = functions_map[tool_call.name](**tool_call.parameters)
6
    # store the output in a list
7
    outputs = [output]
8
    print(f"== tool results: {outputs}")
9
    # store your tool results in this format
10
    tool_results.append({
11
        "call": tool_call,
12
        "outputs": outputs
13
    })
14
15
print("Tool results that will be fed back to the model in step 4:")
16
print(json.dumps(tool_results, indent=4))

1
response = co.chat(
2
    message=message,
3
    tools=tools,
4
    tool_results=tool_results,
5
    preamble=preamble,
6
    model="command-r",
7
    temperature=0.3
8
)

1
print("Final answer:")
2
print(response.text)

1
print("Citations that support the final answer:")
2
for cite in response.citations:
3
  print(cite)

1
def insert_citations_in_order(text, citations):
2
    """
3
    A helper function to pretty print citations.
4
    """
5
    offset = 0
6
    document_id_to_number = {}
7
    citation_number = 0
8
    modified_citations = []
9
10
    # Process citations, assigning numbers based on unique document_ids
11
    for citation in citations:
12
        citation_numbers = []
13
        for document_id in sorted(citation["document_ids"]):
14
            if document_id not in document_id_to_number:
15
                citation_number += 1  # Increment for a new document_id
16
                document_id_to_number[document_id] = citation_number
17
            citation_numbers.append(document_id_to_number[document_id])
18
19
        # Adjust start/end with offset
20
        start, end = citation['start'] + offset, citation['end'] + offset
21
        placeholder = ''.join([f'[{number}]' for number in citation_numbers])
22
        # Bold the cited text and append the placeholder
23
        modification = f'**{text[start:end]}**{placeholder}'
24
        # Replace the cited text with its bolded version + placeholder
25
        text = text[:start] + modification + text[end:]
26
        # Update the offset for subsequent replacements
27
        offset += len(modification) - (end - start)
28
29
    # Prepare citations for listing at the bottom, ensuring unique document_ids are listed once
30
    unique_citations = {number: doc_id for doc_id, number in document_id_to_number.items()}
31
    citation_list = '\n'.join([f'[{doc_id}] source: {tool_results[doc_id - 1]["outputs"]} \n    based on tool call: {dict(tool_results[doc_id - 1]["call"])}' for doc_id, number in sorted(unique_citations.items(), key=lambda item: item[1])])
32
    text_with_citations = f'{text}\n\n{citation_list}'
33
34
    return text_with_citations
35
36
37
print(insert_citations_in_order(response.text, response.citations))