|
|
""" |
|
|
QoScope agent with tools. |
|
|
""" |
|
|
import yaml |
|
|
import sqlalchemy as sqa |
|
|
from Gradio_UI import GradioUI |
|
|
from smolagents import CodeAgent, HfApiModel, tool |
|
|
|
|
|
from tools.diagrams import PlotTool |
|
|
from tools.final_answer import FinalAnswerTool |
|
|
|
|
|
|
|
|
DATABASE_PATH = 'sqlite:///resampled_daily_avg.sqlite' |
|
|
|
|
|
db_engine = sqa.create_engine(DATABASE_PATH, echo=False) |
|
|
|
|
|
|
|
|
@tool |
|
|
def run_sql_query(sql_query: str) -> str: |
|
|
""" |
|
|
Run a SQL query on a table in a SQLite database and return the output/result as a string. |
|
|
The output contains one row of result(s) in each line. |
|
|
The output is simple text without any Markdown styling. An agent should take this plain output |
|
|
and format appropriately when required, e.g., as a Markdown table or summarizing the results |
|
|
in words. |
|
|
|
|
|
The only available table in the SQLite database is `school_measurements`. |
|
|
The table's CREATE statement (DDL) is as follows: |
|
|
CREATE TABLE school_measurements ( |
|
|
"index" BIGINT, |
|
|
date DATETIME, |
|
|
download_speed FLOAT, |
|
|
upload_speed FLOAT, |
|
|
latency FLOAT, |
|
|
school_id_giga TEXT, |
|
|
school_name TEXT, |
|
|
server_location TEXT, |
|
|
country TEXT, |
|
|
iso3_format TEXT |
|
|
) |
|
|
|
|
|
IMPORTANT: You will ONLY execute SELECT queries. You will NEVER execute any other types of |
|
|
SQL queries, e.g., INSERT, DELETE, DROP, and so on, which changes the database in any way. |
|
|
|
|
|
Args: |
|
|
sql_query: An appropriate, correct SQL query for a SQLite database |
|
|
|
|
|
Returns: |
|
|
The result of running the SQL query as a string. |
|
|
""" |
|
|
|
|
|
|
|
|
|
|
|
output = '' |
|
|
with db_engine.connect() as con: |
|
|
rows = con.execute(sqa.text(sql_query)) |
|
|
for row in rows: |
|
|
|
|
|
output += '\n' + str(row) |
|
|
|
|
|
return output |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
final_answer = FinalAnswerTool() |
|
|
code_model = HfApiModel( |
|
|
max_tokens=2096, |
|
|
temperature=0.2, |
|
|
model_id='Qwen/Qwen2.5-Coder-32B-Instruct', |
|
|
custom_role_conversions=None, |
|
|
) |
|
|
|
|
|
with open('prompts.yaml', 'r', encoding='utf-8') as stream: |
|
|
prompt_templates = yaml.safe_load(stream) |
|
|
|
|
|
agent = CodeAgent( |
|
|
model=code_model, |
|
|
tools=[ |
|
|
run_sql_query, |
|
|
PlotTool(), |
|
|
final_answer, |
|
|
], |
|
|
max_steps=6, |
|
|
verbosity_level=1, |
|
|
grammar=None, |
|
|
planning_interval=None, |
|
|
name=None, |
|
|
description=None, |
|
|
prompt_templates=prompt_templates |
|
|
) |
|
|
|
|
|
GradioUI(agent).launch() |
|
|
|
