Querying Structured Data (SQL)

1
! pip install cohere pandas -qq

1
import json
2
import os
3
import cohere
4
import sqlite3
5
import pandas as pd
6
7
co = cohere.ClientV2(
8
    "COHERE_API_KEY"
9
)  # Get your free API key: https://dashboard.cohere.com/api-keys

1
# Create a connection to a new SQLite database (or connect to an existing one)
2
conn = sqlite3.connect("evaluation_results.db")
3
cursor = conn.cursor()
4
5
# Execute the CREATE TABLE command
6
cursor.execute(
7
    """
8
CREATE TABLE evaluation_results (
9
    usecase TEXT,
10
    run TEXT,
11
    score FLOAT,
12
    temperature FLOAT,
13
    tokens INTEGER,
14
    latency FLOAT
15
)
16
"""
17
)
18
19
# Execute the INSERT commands
20
data = [
21
    ("extract_names", "A", 0.5, 0.3, 103, 1.12),
22
    ("draft_email", "A", 0.6, 0.3, 252, 2.5),
23
    ("summarize_article", "A", 0.8, 0.3, 350, 4.2),
24
    ("extract_names", "B", 0.2, 0.3, 101, 2.85),
25
    ("draft_email", "B", 0.4, 0.3, 230, 3.2),
26
    ("summarize_article", "B", 0.6, 0.3, 370, 4.2),
27
    ("extract_names", "C", 0.7, 0.3, 101, 2.22),
28
    ("draft_email", "C", 0.5, 0.3, 221, 2.5),
29
    ("summarize_article", "C", 0.1, 0.3, 361, 3.9),
30
    ("extract_names", "D", 0.7, 0.5, 120, 3.2),
31
    ("draft_email", "D", 0.8, 0.5, 280, 3.4),
32
    ("summarize_article", "D", 0.9, 0.5, 342, 4.8),
33
]
34
35
cursor.executemany(
36
    "INSERT INTO evaluation_results VALUES (?,?,?,?,?,?)", data
37
)
38
39
# Commit the changes and close the connection
40
conn.commit()
41
conn.close()

1
def sql_table_query(query: str) -> dict:
2
    """
3
    Execute an SQL query on the evaluation_results table and return the result as a dictionary.
4
5
    Args:
6
    query (str): SQL query to execute on the evaluation_results table
7
8
    Returns:
9
    dict: Result of the SQL query
10
    """
11
    try:
12
        # Connect to the SQLite database
13
        conn = sqlite3.connect("evaluation_results.db")
14
15
        # Execute the query and fetch the results into a DataFrame
16
        df = pd.read_sql_query(query, conn)
17
18
        # Close the connection
19
        conn.close()
20
21
        # Convert DataFrame to dictionary
22
        result_dict = df.to_dict(orient="records")
23
24
        return result_dict
25
26
    except sqlite3.Error as e:
27
        print(f"An error occurred: {e}")
28
        return str(e)
29
    except Exception as e:
30
        print(f"An unexpected error occurred: {e}")
31
        return str(e)
32
33
34
functions_map = {"sql_table_query": sql_table_query}

1
result = sql_table_query(
2
    "SELECT * FROM evaluation_results WHERE usecase = 'extract_names'"
3
)
4
print(result)

1
sql_table_query_tool = {
2
    "type": "function",
3
    "function": {
4
        "name": "sql_table_query",
5
        "description": "Execute an SQL query on the evaluation_results table in the SQLite database. The table has columns 'usecase', 'run', 'score', 'temperature', 'tokens', and 'latency'.",
6
        "parameters": {
7
            "type": "object",
8
            "properties": {
9
                "query": {
10
                    "type": "string",
11
                    "description": "SQL query to execute on the evaluation_results table",
12
                }
13
            },
14
            "required": ["query"],
15
        },
16
    },
17
}
18
19
tools = [sql_table_query_tool]

1
system_message = """## Task and Context
2
You are an assistant who helps developers analyze LLM application evaluation results from a SQLite database. The database contains a table named 'evaluation_results' with the following schema:
3
4
- usecase (TEXT): The type of task being evaluated
5
- run (TEXT): The identifier for a specific evaluation run
6
- score (REAL): The performance score of the run
7
- temperature (REAL): The temperature setting used for the LLM
8
- tokens (INTEGER): The number of tokens used in the run
9
- latency (REAL): The time taken for the run in seconds
10
11
You can use SQL queries to analyze this data and provide insights to the developers."""

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