afa-bd_sql-service / query_engine.py
pdrMottaS
creating service
d128798
raw
history blame
942 Bytes
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():
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("duckdb:///llm_view.db")
sql_database = SQLDatabase(engine)
return NLSQLTableQueryEngine(
sql_database=sql_database, # The SQL database instance to query
tables=["afabd"], # List of tables to include in the query engine
llm=llm, # The language model used for processing natural language queries
)