Querying Structured Data (Tables)

1
import json
2
import os
3
import cohere
4
5
co = cohere.ClientV2("COHERE_API_KEY") # Get your free API key: https://dashboard.cohere.com/api-keys

1
import pandas as pd
2
3
df = pd.read_csv('evaluation_results.csv')
4
5
df.head()

1
from tool_def import (
2
    analyze_evaluation_results,
3
    analyze_evaluation_results_tool,
4
)

1
functions_map = {
2
    "analyze_evaluation_results": analyze_evaluation_results
3
}

1
analyze_evaluation_results_tool = {
2
        "type": "function",
3
        "function": {
4
            "name": "analyze_evaluation_results",
5
            "description": "Generate Python code using the pandas library to analyze evaluation results from a dataframe called `evaluation_results`. The dataframe has columns 'usecase','run','score','temperature','tokens', and 'latency'. You must start with `import pandas as pd` and read a CSV file called `evaluation_results.csv` into the `evaluation_results` dataframe.",
6
            "parameters": {
7
                "type": "object",
8
                "properties": {
9
                    "code": {
10
                        "type": "string",
11
                        "description": "Executable Python code"
12
                    }
13
                },
14
                "required": ["code"]
15
            }
16
        }
17
}

1
tools = [analyze_evaluation_results_tool]

1
system_message="""## Task and Context
2
You are an assistant who helps developers find code examples and tutorials on using Cohere."""

1
model = "command-r-plus-08-2024"
2
3
def run_agent(query, messages=None):
4
    if messages is None:
5
        messages = []
6
        
7
    if "system" not in {m.get("role") for m in messages}:
8
        messages.append({"role": "system", "content": system_message})
9
    
10
    # Step 1: get user message
11
    print(f"Question:\n{query}")
12
    print("="*50)
13
    
14
    messages.append({"role": "user", "content": query})
15
16
    # Step 2: Generate tool calls (if any)
17
    response = co.chat(
18
        model=model,
19
        messages=messages,
20
        tools=tools,
21
        temperature=0.1
22
    )
23
24
    while response.message.tool_calls:
25
        
26
        print("TOOL PLAN:")
27
        print(response.message.tool_plan,"\n")
28
        print("TOOL CALLS:")
29
        for tc in response.message.tool_calls:
30
            if tc.function.name == "analyze_evaluation_results":
31
                print(f"Tool name: {tc.function.name}")
32
                tool_call_prettified = print("\n".join(f"  {line}" for line_num, line in enumerate(json.loads(tc.function.arguments)["code"].splitlines())))
33
                print(tool_call_prettified)
34
            else:
35
                print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
36
        print("="*50)
37
38
        messages.append({"role": "assistant", "tool_calls": response.message.tool_calls, "tool_plan": response.message.tool_plan})        
39
        
40
        # Step 3: Get tool results
41
        for tc in response.message.tool_calls:
42
            tool_result = functions_map[tc.function.name](**json.loads(tc.function.arguments))
43
            tool_content = [({"type": "document", "document": {"data": json.dumps(tool_result)}})]
44
                
45
            messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_content})
46
        
47
        # Step 4: Generate response and citations 
48
        response = co.chat(
49
            model=model,
50
            messages=messages,
51
            tools=tools,
52
            temperature=0.1
53
        )
54
    
55
    messages.append({"role": "assistant", "content": response.message.content[0].text})
56
        
57
    # Print final response
58
    print("RESPONSE:")
59
    print(response.message.content[0].text)
60
    print("="*50)
61
    
62
    # Print citations (if any)
63
    verbose_source = False # Change to True to display the contents of a source
64
    if response.message.citations:
65
        print("CITATIONS:\n")
66
        for citation in response.message.citations:
67
            print(f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' ")
68
            print("Sources:")
69
            for idx, source in enumerate(citation.sources):
70
                print(f"{idx+1}. {source.id}")
71
                if verbose_source:
72
                    print(f"{source.tool_output}")
73
            print("\n")
74
                    
75
    return messages

1
messages = run_agent("What's the average evaluation score in run A")
2
# Answer: 0.63

1
messages = run_agent("What's the latency of the highest-scoring run for the summarize_article use case?")
2
# Answer: 4.8

1
messages = run_agent("Which use case uses the least amount of tokens on average? Show the comparison of all use cases in a markdown table.")
2
# Answer: extract_names (106.25), draft_email (245.75), summarize_article (355.75)