from llama_index.llms.groq import Groq
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
from sqlalchemy import create_engine
from llama_index.core import SQLDatabase
from llama_index.core.indices.struct_store import NLSQLTableQueryEngine

def set_query_engine(path:str):
    llm = Groq(
        model="llama3-8b-8192",
        api_key="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"
    )
    embedding = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
    Settings.llm = llm
    Settings.embed_model = embedding
    engine = create_engine(f"duckdb:///{path}")
    sql_database = SQLDatabase(engine)
    return NLSQLTableQueryEngine(
        sql_database=sql_database, # The SQL database instance to query
        tables=["escola", "curso", "avaliacao"], # List of tables to include in the query engine
        llm=llm, # The language model used for processing natural language queries
    )