Spaces:
Running
Running
| 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 | |