manabcommercial / AskQuestion.py
manabb's picture
Update AskQuestion.py
39b5dfe verified
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
# from langchain_classic.chains import create_sql_query_chain
from langchain.chains import create_sql_query_chain
from langchain_core.prompts import PromptTemplate
import os
def ask_question(question: str):
# Connect to your DB (same as FastAPI)
db = SQLDatabase.from_uri("sqlite:///./manabCommercial1.db")
# db = SQLDatabase.from_uri("sqlite:////data/manabCommercial1.db")
# LLM (use your API key)
api_key = os.getenv("OPENAI_API_KEY")
if api_key:
os.environ["OPENAI_API_KEY"] = api_key
llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0
)
table_info = db.get_table_info()
print("Table Info:", table_info) # Debugging line to check schema retrieval
top_k = 5 # Limit to top 5 rows
input=question
# Custom prompt for your schema
PROMPT = """
You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run.
Schema: {table_info}
Only use: SELECT from manabCommercial1. No INSERT/UPDATE/DELETE.
Limit to top {top_k} rows.
Question: {input}
SQLQuery:
"""
prompt = PromptTemplate.from_template(PROMPT)
# Create chain
chain = create_sql_query_chain(llm, db, prompt, top_k=top_k)
# Query example
# question = "Show status for EFileNo containing 'ABC' where completed=Yes"
sql_query = chain.invoke({"input": question})
print(sql_query) # "SELECT status FROM manabCommercial1 WHERE EFileNo LIKE '%ABC%' AND completed='Yes'"
# Execute + Answer
result = db.run(sql_query)
full_response = llm.invoke(f"Query: {sql_query}\nResult: {result}\nAnswer the original question: {question}")
# print(full_response.content)
return full_response.content