Update app.py

app.py

@@ -9,7 +9,10 @@ from groq import Groq
 # Fetch API key from environment variable
 API_KEY = os.environ.get('GroqApi')
 
-#API_KEY = os.environ.get("GROQ_API_KEY")
+# Initialize variables
+INDEX = None
+CHUNKS = None
+MODEL = None
 
 # Function to scrape tariff data
 def scrape_tariff_data(url):
@@ -37,17 +40,25 @@ def create_faiss_index(chunks, model_name='all-MiniLM-L6-v2'):
     index.add(embeddings)
     return index, embeddings, model
 
-# Function to query the Groq API
-def query_llm(prompt):
+# Function to search FAISS for relevant chunks
+def search_faiss(query, index, chunks, model, top_k=5):
+    query_embedding = model.encode([query])
+    distances, indices = index.search(query_embedding, top_k)
+    relevant_chunks = [chunks[i] for i in indices[0] if i < len(chunks)]
+    return relevant_chunks
+
+# Function to query the Groq API with augmented query
+def query_llm(prompt, context):
     if not API_KEY:
         return "Error: GROQ_API_KEY is not set in environment variables."
     
     client = Groq(api_key=API_KEY)
+    augmented_prompt = f"Based on the following data:\n\n{context}\n\nAnswer the question: {prompt}"
     chat_completion = client.chat.completions.create(
         messages=[
             {
                 "role": "user",
-                "content": prompt,
+                "content": augmented_prompt,
             }
         ],
         model="llama3-8b-8192",
@@ -63,8 +74,9 @@ if st.button("Process Tariff Data"):
     with st.spinner("Extracting and processing data..."):
         try:
             text = scrape_tariff_data(url)
-            chunks = chunk_text(text)
-            index, embeddings, model = create_faiss_index(chunks)
+            global CHUNKS, INDEX, MODEL
+            CHUNKS = chunk_text(text)
+            INDEX, embeddings, MODEL = create_faiss_index(CHUNKS)
             st.success("Data processed and indexed!")
         except Exception as e:
             st.error(f"Error processing data: {e}")
@@ -76,8 +88,16 @@ if st.button("Get Answer"):
     if prompt:
         with st.spinner("Fetching response..."):
             try:
-                response = query_llm(prompt)
-                st.write(response)
+                if not (INDEX and CHUNKS and MODEL):
+                    st.error("Data has not been processed yet. Please process the data first.")
+                else:
+                    # Retrieve relevant chunks
+                    relevant_chunks = search_faiss(prompt, INDEX, CHUNKS, MODEL)
+                    context = "\n".join(relevant_chunks)
+
+                    # Query the LLM with context
+                    response = query_llm(prompt, context)
+                    st.write(response)
             except Exception as e:
                 st.error(f"Error querying the model: {e}")
     else: