Update search.py

search.py

@@ -1,23 +1,26 @@
 import os
 import streamlit as st
 from langchain_groq import ChatGroq
-from langchain_community.tools import ArxivQueryRun,WikipediaQueryRun,DuckDuckGoSearchRun
-from langchain_community.utilities import WikipediaAPIWrapper,ArxivAPIWrapper
+from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun, DuckDuckGoSearchRun
+from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
 from langchain.agents import initialize_agent, AgentType
 from langchain.callbacks import StreamlitCallbackHandler
 from dotenv import load_dotenv
 
-##CODE
+# Load environment variables from .env file
+load_dotenv()  
+api_key = os.getenv("GROQ_API_KEY")  # Get the API key from the environment variables
+
 ## Arxiv and Wikipedia Tools
-arxiv_wrapper=ArxivAPIWrapper(top_k_results=1, doc_content_chars_max=200)
-arxiv=ArxivQueryRun(api_wrapper=arxiv_wrapper)
+arxiv_wrapper = ArxivAPIWrapper(top_k_results=1, doc_content_chars_max=200)
+arxiv = ArxivQueryRun(api_wrapper=arxiv_wrapper)
 
-api_wrapper=WikipediaAPIWrapper(top_k_results=1,doc_content_chars_max=200)
-wiki=WikipediaQueryRun(api_wrapper=api_wrapper)
+api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=200)
+wiki = WikipediaQueryRun(api_wrapper=api_wrapper)
 
-search=DuckDuckGoSearchRun(name="Search")
+search = DuckDuckGoSearchRun(name="Search")
 
-##NILESH
+## NILESH
 st.title("🤖 NileAI - Your AI Search Companion")
 """
 Welcome to NileSearch, your AI-powered chat agent for real-time web search and insights.
@@ -25,33 +28,29 @@ This app uses `StreamlitCallbackHandler` to display the agent's thoughts and act
 Explore more LangChain 🧠 Streamlit Agent examples at [github.com/langchain-ai/streamlit-agent](https://github.com/langchain-ai/streamlit-agent).
 """
 
-
 ## Sidebar for Settings
 # st.sidebar.title("Settings")
 # api_key = st.sidebar.text_input("Enter your Groq API key:", type="password")
 
-api_key="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"
-
 if "messages" not in st.session_state:
     st.session_state["messages"] = [
         {"role": "assistant", "content": "Hi! How can I assist you today?"}
     ]
 
-    
 for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
-    
-if prompt:=st.chat_input(placeholder="What is machine learning?"):
-    st.session_state.messages.append({"role":"user", "content":prompt})
+
+if prompt := st.chat_input(placeholder="What is machine learning?"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
-    
+
     llm = ChatGroq(groq_api_key=api_key, model_name="Llama3-8b-8192", streaming=True)
     tools = [search, arxiv, wiki]
-    
-    search_agent = initialize_agent(tools, llm, agent = AgentType.ZERO_SHOT_REACT_DESCRIPTION, handle_parsing_errors=True)
-    
+
+    search_agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, handle_parsing_errors=True)
+
     with st.chat_message("assistant"):
         st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=False)
         response = search_agent.run(st.session_state.messages, callbacks=[st_cb])
-        st.session_state.messages.append({'role':'assistant', 'content':response})
+        st.session_state.messages.append({'role': 'assistant', 'content': response})
         st.write(response)