updated with langchain sql agent

app.py

@@ -1,19 +1,15 @@
 import pandas as pd
-import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig,
-    pipeline
-)
-from langchain_experimental.agents import create_pandas_dataframe_agent
-from langchain_community.llms import HuggingFacePipeline
+import os
+from dotenv import load_dotenv
+from langchain.chat_models import init_chat_model
+from langchain_community.agent_toolkits import create_sql_agent
+from langchain_community.utilities import SQLDatabase
+from sqlalchemy import create_engine
 import gradio as gr
-import spaces
-
+load_dotenv()
 
+llm = init_chat_model("gemini-2.5-flash", model_provider="google_genai")
 # --- Config ---
-LLM_MODEL_ID = "HuggingFaceH4/zephyr-7b-beta"
 DATA_FILE = "IPL.csv"
 
 
@@ -32,86 +28,18 @@ def load_df():
 
 _df = load_df()
 
-
-# --- Load Quantized Model ---
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=False,
-)
-
-tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
-
-model = AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    quantization_config=bnb_config,
-    trust_remote_code=True,
-)
-
-
-# --- LLM Wrapper for LangChain ---
-class MyLLMWrapper:
-    def __init__(self):
-        self.tokenizer = tokenizer
-        self.model = model
-
-    def invoke(self, input_str):
-        return self.__call__(input_str)
-
-    def __call__(self, input_str):
-        inputs = self.tokenizer(input_str, return_tensors="pt").to(self.model.device)
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=256,
-                do_sample=True,
-                temperature=0.1,
-                top_p=0.9,
-                eos_token_id=self.tokenizer.eos_token_id,
-                pad_token_id=self.tokenizer.pad_token_id
-            )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-
-
-llm = MyLLMWrapper()
-
-
-# --- System Prompt for the Agent ---
-system_message = """
-You are an expert IPL cricket analyst. You have access to a pandas DataFrame named `df` that contains ball-by-ball IPL match data.
-Answer all questions using pandas logic, match stats, and accurate calculations.
-"""
-
-# --- LangChain Agent ---
-agent = create_pandas_dataframe_agent(
-    llm,
-    _df,
-    verbose=False,
-    max_execution_time=None,
-    early_stopping_method="force",
-    include_df_in_prompt=True,
-    number_of_head_rows=5,
-    extra_tools=(),
-    # handle_parsing_errors=True,
-    agent_executor_kwargs={"system_message": system_message},
-    agent_type="openai-tools",  # Most compatible with Hugging Face models
-    allow_dangerous_code=True,
-)
+engine = create_engine("sqlite:///ipl.db")
+_df.to_sql("ipl", engine, index=False)
+db = SQLDatabase(engine=engine)
+print("Db created successfully")
 
 
-# --- Inference Function ---
-@spaces.GPU(duration=120)
-def predict_answer(question):
-    torch.cuda.empty_cache()
+def main(query):
     try:
-        res = agent.invoke({"input": question})
-        return res.get("output", "No response generated.")
-    except Exception as e:
-        return f"❌ Error during inference: {e}"
+        agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
+        return agent_executor.invoke({"input": query})
+    except:
+        return "Failed to fetch the required info"
 
 
 # --- Gradio UI ---
@@ -133,7 +61,7 @@ with gr.Blocks() as demo:
 
     def bot_reply(hist):
         q = hist[-1][0]
-        a = predict_answer(q)
+        a = main(q)
         hist[-1][1] = a
         return hist