Update interim.py

interim.py

@@ -8,6 +8,7 @@ from datetime import datetime, timezone
 from crewai import Agent, Crew, Process, Task
 from crewai.tools import tool
 from langchain_groq import ChatGroq
+from langchain_openai import ChatOpenAI
 from langchain.schema.output import LLMResult
 from langchain_core.callbacks.base import BaseCallbackHandler
 from langchain_community.tools.sql_database.tool import (
@@ -23,30 +24,33 @@ 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=120,
-    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
-)
-
 st.title("Blah Blah App 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
@@ -148,7 +152,7 @@ if st.session_state.df is not None:
         verbose=True,
     )
 
-    # Query Input for Patent Analysis
+    #Query Input for Patent Analysis
     query = st.text_area("Enter Patent Analysis Query:", placeholder="e.g., 'How many patents related to Machine Learning were filed after 2016?'")
     if st.button("Submit Query"):
         with st.spinner("Processing your query..."):