Update app.py

app.py

@@ -20,33 +20,33 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
-# API Key
-os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
-
-# Initialize LLM
-class LLMCallbackHandler(BaseCallbackHandler):
-    def __init__(self, log_path: Path):
-        self.log_path = log_path
-
-    def on_llm_start(self, serialized, prompts, **kwargs):
-        with self.log_path.open("a", encoding="utf-8") as file:
-            file.write(json.dumps({"event": "llm_start", "text": prompts[0], "timestamp": datetime.now().isoformat()}) + "\n")
-
-    def on_llm_end(self, response: LLMResult, **kwargs):
-        generation = response.generations[-1][-1].message.content
-        with self.log_path.open("a", encoding="utf-8") as file:
-            file.write(json.dumps({"event": "llm_end", "text": generation, "timestamp": datetime.now().isoformat()}) + "\n")
-
-llm = ChatGroq(
-    temperature=0,
-    model_name="groq/llama-3.3-70b-versatile",
-    max_tokens=500,
-    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
-)
-
 st.title("SQL-RAG Using CrewAI 🚀")
 st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
 
+# Initialize LLM
+llm = None
+
+# Model Selection
+model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
+
+
+# API Key Validation and LLM Initialization
+groq_api_key = os.getenv("GROQ_API_KEY")
+openai_api_key = os.getenv("OPENAI_API_KEY")
+
+if model_choice == "llama-3.3-70b":
+    if not groq_api_key:
+        st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
+        llm = None
+    else:
+        llm = ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
+elif model_choice == "GPT-4o":
+    if not openai_api_key:
+        st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
+        llm = None
+    else:
+        llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
+
 # Initialize session state for data persistence
 if "df" not in st.session_state:
     st.session_state.df = None