Querying Structured Data (SQL)

1
import json
2
import os
3
import cohere
4
import sqlite3
5
import pandas as pd
6
7
co = cohere.ClientV2("COHERE_API_KEY") # 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
CREATE TABLE evaluation_results (
8
    usecase TEXT,
9
    run TEXT,
10
    score FLOAT,
11
    temperature FLOAT,
12
    tokens INTEGER,
13
    latency FLOAT
14
)
15
''')
16
17
# Execute the INSERT commands
18
data = [
19
    ('extract_names', 'A', 0.5, 0.3, 103, 1.12),
20
    ('draft_email', 'A', 0.6, 0.3, 252, 2.5),
21
    ('summarize_article', 'A', 0.8, 0.3, 350, 4.2),
22
    ('extract_names', 'B', 0.2, 0.3, 101, 2.85),
23
    ('draft_email', 'B', 0.4, 0.3, 230, 3.2),
24
    ('summarize_article', 'B', 0.6, 0.3, 370, 4.2),
25
    ('extract_names', 'C', 0.7, 0.3, 101, 2.22),
26
    ('draft_email', 'C', 0.5, 0.3, 221, 2.5),
27
    ('summarize_article', 'C', 0.1, 0.3, 361, 3.9),
28
    ('extract_names', 'D', 0.7, 0.5, 120, 3.2),
29
    ('draft_email', 'D', 0.8, 0.5, 280, 3.4),
30
    ('summarize_article', 'D', 0.9, 0.5, 342, 4.8)
31
]
32
33
cursor.executemany('INSERT INTO evaluation_results VALUES (?,?,?,?,?,?)', data)
34
35
# Commit the changes and close the connection
36
conn.commit()
37
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
functions_map = {
34
    "sql_table_query": sql_table_query
35
}

1
result = sql_table_query("SELECT * FROM evaluation_results WHERE usecase = 'extract_names'")
2
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-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
            # print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
31
            if tc.function.name == "analyze_evaluation_results":
32
                print(f"Tool name: {tc.function.name}")
33
                tool_call_prettified = print("\n".join(f"  {line}" for line_num, line in enumerate(json.loads(tc.function.arguments)["code"].splitlines())))
34
                print(tool_call_prettified)
35
            else:
36
                print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
37
        print("="*50)
38
39
        messages.append({"role": "assistant", "tool_calls": response.message.tool_calls, "tool_plan": response.message.tool_plan})        
40
        
41
        # Step 3: Get tool results
42
        for tc in response.message.tool_calls:
43
            tool_result = functions_map[tc.function.name](**json.loads(tc.function.arguments))
44
            tool_content = [({"type": "document", "document": {"data": json.dumps(tool_result)}})]
45
                
46
            messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_content})
47
        
48
        # Step 4: Generate response and citations 
49
        response = co.chat(
50
            model=model,
51
            messages=messages,
52
            tools=tools,
53
            temperature=0.1
54
        )
55
    
56
    messages.append({"role": "assistant", "content": response.message.content[0].text})
57
        
58
    # Print final response
59
    print("Response:")
60
    print(response.message.content[0].text)
61
    print("="*50)
62
    
63
    # Print citations (if any)
64
    verbose_source = False # Change to True to display the contents of a source
65
    if response.message.citations:
66
        print("CITATIONS:\n")
67
        for citation in response.message.citations:
68
            print(f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' ")
69
            print("Sources:")
70
            for idx, source in enumerate(citation.sources):
71
                print(f"{idx+1}. {source.id}")
72
                if verbose_source:
73
                    print(f"{source.tool_output}")
74
            print("\n")
75
    
76
    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)