Querying Structured Data (Tables)

1
! pip install cohere pandas -qq

1
import json
2
import os
3
import cohere
4
5
co = cohere.ClientV2(
6
    "COHERE_API_KEY"
7
)  # 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
ou are an assistant who helps developers analyze LLM application evaluation results from a CSV files."""

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
            if tc.function.name == "analyze_evaluation_results":
29
                print(f"Tool name: {tc.function.name}")
30
                tool_call_prettified = print(
31
                    "\n".join(
32
                        f"  {line}"
33
                        for line_num, line in enumerate(
34
                            json.loads(tc.function.arguments)[
35
                                "code"
36
                            ].splitlines()
37
                        )
38
                    )
39
                )
40
                print(tool_call_prettified)
41
            else:
42
                print(
43
                    f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}"
44
                )
45
        print("=" * 50)
46
47
        messages.append(
48
            {
49
                "role": "assistant",
50
                "tool_calls": response.message.tool_calls,
51
                "tool_plan": response.message.tool_plan,
52
            }
53
        )
54
55
        # Step 3: Get tool results
56
        for tc in response.message.tool_calls:
57
            tool_result = functions_map[tc.function.name](
58
                **json.loads(tc.function.arguments)
59
            )
60
            tool_content = [
61
                (
62
                    {
63
                        "type": "document",
64
                        "document": {"data": json.dumps(tool_result)},
65
                    }
66
                )
67
            ]
68
69
            messages.append(
70
                {
71
                    "role": "tool",
72
                    "tool_call_id": tc.id,
73
                    "content": tool_content,
74
                }
75
            )
76
77
        # Step 4: Generate response and citations
78
        response = co.chat(
79
            model=model,
80
            messages=messages,
81
            tools=tools,
82
            temperature=0.3,
83
        )
84
85
    messages.append(
86
        {
87
            "role": "assistant",
88
            "content": response.message.content[0].text,
89
        }
90
    )
91
92
    # Print final response
93
    print("RESPONSE:")
94
    print(response.message.content[0].text)
95
    print("=" * 50)
96
97
    # Print citations (if any)
98
    verbose_source = (
99
        False  # Change to True to display the contents of a source
100
    )
101
    if response.message.citations:
102
        print("CITATIONS:\n")
103
        for citation in response.message.citations:
104
            print(
105
                f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' "
106
            )
107
            print("Sources:")
108
            for idx, source in enumerate(citation.sources):
109
                print(f"{idx+1}. {source.id}")
110
                if verbose_source:
111
                    print(f"{source.tool_output}")
112
            print("\n")
113
114
    return messages

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

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

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