Spaces:
Configuration error
Configuration error
File size: 2,748 Bytes
2fccb82 f959207 c92e861 2fccb82 c92e861 2fccb82 f959207 2fccb82 f959207 2fccb82 |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import mysql.connector
import os
from dotenv import load_dotenv
from smolagents import CodeAgent, HfApiModel, tool
load_dotenv()
connection = None
def set_connection(conn):
"""
Setter global untuk connection. Connection ini bisa diubah dari file lain.
"""
global connection
connection = conn
def createConnection(user='root',password='password',host='127.0.0.1',port='3306',database='sakila'):
cnx = mysql.connector.connect(user=user, password=password,
host=host,
database=database,
port=port)
return cnx
def getAllTable(connection):
cursor = connection.cursor()
cursor.execute("SHOW TABLES")
tables = cursor.fetchall()
return tables
def createColumnDesc(connection):
DESC_TABLE = """"""
cursor = connection.cursor()
tables = getAllTable(connection)
for table in tables:
cursor.execute(f"DESCRIBE {table[0]}")
columns = cursor.fetchall()
DESC_TABLE += f"*{table[0].upper()}\n"
for column in columns:
column_name = column[0]
column_type = column[1]
DESC_TABLE += f"- {column_name} : {column_type}\n"
return DESC_TABLE
@tool
def sql_engine(query: str) -> str:
"""
Performs SQL queries and returns results as a minimal string for LLM processing.
Args:
query: The SQL query to execute. This should be valid SQL.
Returns:
str: Query results in a concise, consistent, and minimal format.
"""
cursor = connection.cursor(dictionary=True) # Enable dictionary cursor
cursor.execute(query)
# Get column names
columns = [desc[0] for desc in cursor.description]
# Fetch all rows
rows = cursor.fetchall()
# Return results in a flat format
if not rows:
return "No results found"
result = []
result.append("|".join(columns)) # Column headers
for row in rows:
result.append("|".join(str(row[col]) for col in columns))
return "\n".join(result)
token = os.getenv('HUGGINGFACE_TOKEN')
agent = CodeAgent(
tools=[sql_engine],
model=HfApiModel("Qwen/Qwen2.5-Coder-32B-Instruct",token=token),
additional_authorized_imports=["matplotlib.pyplot"]
)
def Inference(conn,query):
connection = conn
desc = createColumnDesc(connection)
DESCRIPTION = """
Allows you to perform SQL queries on the table and returns results in various formats. Beware that this tool's output is a string representation of the execution output.
It can use the following tables:
"""
DESCRIPTION += desc
sql_engine.description = DESCRIPTION
result = agent.run(query)
return str(result)
|