File size: 2,492 Bytes
fb44830 4e29b9c fb44830 4e29b9c fb44830 4e29b9c fb44830 4e29b9c fb44830 4e29b9c fb44830 4e29b9c fb44830 4e29b9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
"""
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.
"""
# Source: https://huggingface.co/docs/smolagents/en/examples/text_to_sql
output = ''
with db_engine.connect() as con:
rows = con.execute(sqa.text(sql_query))
for row in rows:
# Each row is a tuple
output += '\n' + str(row)
return output
### Main block ###
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()
|