File size: 2,023 Bytes
c976bfb |
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 |
import os
from sqlalchemy.sql import text
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from semantic_kernel import Kernel
from semantic_kernel.orchestration.sk_context import SKContext
from semantic_kernel.skill_definition import sk_function, sk_function_context_parameter
# Load environment variables from .env file
from dotenv import load_dotenv
load_dotenv()
class SQLGetQuery:
def __init__(self, kernel: Kernel):
self._kernel = kernel
@sk_function(
description="Query SQL database related to user question",
name="dbQuery",
)
@sk_function_context_parameter(name="input", description="SQL query related to user question which was generated by LLM")
def db_query(self, context: SKContext) -> str:
# Set SQLAlchemy logging level to WARNING
username = os.environ['USERNAME']
password = os.environ['PASSWORD']
driver = os.environ['DRIVER']
# schema = context['db_schema']
server = os.environ['CONNECTION_STRING']
database = os.environ['DB']
query = text(context['input'])
print(f"\n------------------\nSQL QUERY:\n {query}\n------------------\n")
connection_string = f'mssql+pyodbc://{username}:{password}@{server}/{database}?driver={driver}&TrustServerCertificate=no&Authentication=ActiveDirectoryPassword'
# Start query
engine = create_engine(connection_string, echo=False, connect_args={'auto_commit': True}, fast_executemany=True)
Session = sessionmaker(bind=engine)
session = Session()
# Execute the query and fetch the results
results = session.execute(query).fetchall()
results = str(results).strip()
#Update context
context['queryresult'] = results
context['db_string'] = server
context['db_name'] = database
context['sql_query'] = str(query)
print(f"\n------------------\nSQL RESULT:\n {results}\n------------------\n")
return results |