Update app.py

app.py

@@ -5,10 +5,17 @@ import os
 from google.oauth2 import service_account
 from googleapiclient.discovery import build
 from streamlit_chat import message as st_message
-from langchain.schema import SystemMessage
+import plotly.express as px
+from langchain.schema import HumanMessage, SystemMessage, AIMessage
+from langchain.chat_models import ChatOpenAI
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.prompts import PromptTemplate
+import warnings
+import time
 from langchain_groq import ChatGroq
+import numpy as np
 from dotenv import load_dotenv
-import warnings
+import re
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -31,18 +38,27 @@ Instructions:
 3. Provide information in a single sentence or a short, structured response.
 4. If the requested information isn’t available or verifiable, respond with "Information not available."
 
+#### Example Output Format:
+"Company: {entity} | Requested Information: {extracted_information}"
+
 Begin extraction.
 """
 
+# Function to get response from LLM
 def get_llm_response(entity, query):
-    formatted_prompt = PROMPT_TEMPLATE.format(entity=entity, query=query)
-    response = llm([SystemMessage(content=formatted_prompt)])
+    # Format the prompt with the entity and query
+    prompt = PROMPT_TEMPLATE.format(entity=entity, query=query)
+    
+    # Request response from the LLM
+    response = llm([SystemMessage(content=prompt)])
+    
+    # Return content or default message
     return response[0].content if response else "Information not available"
 
-# Streamlit Setup
+# Streamlit app setup
 st.set_page_config(page_title="DataScribe", page_icon=":notebook_with_decorative_cover:", layout="wide")
 
-# Sidebar Navigation
+# Sidebar navigation
 with st.sidebar:
     selected = option_menu(
         "DataScribe Menu",
@@ -63,9 +79,47 @@ if "results" not in st.session_state:
 if "column_selection" not in st.session_state:
     st.session_state["column_selection"] = None
 
+# Helper function for Google Sheets API setup
+def get_google_sheet_data(sheet_id, range_name):
+    credentials = service_account.Credentials.from_service_account_info(st.secrets["gcp_service_account"])
+    service = build('sheets', 'v4', credentials=credentials)
+    sheet = service.spreadsheets()
+    result = sheet.values().get(spreadsheetId=sheet_id, range=range_name).execute()
+    values = result.get('values', [])
+    return pd.DataFrame(values[1:], columns=values[0])
+
+# Function to write results back to Google Sheets
+def update_google_sheet(sheet_id, range_name, data):
+    credentials = service_account.Credentials.from_service_account_info(st.secrets["gcp_service_account"])
+    service = build('sheets', 'v4', credentials=credentials)
+    sheet = service.spreadsheets()
+    body = {
+        'values': [data.columns.tolist()] + data.values.tolist()
+    }
+    sheet.values().update(
+        spreadsheetId=sheet_id,
+        range=range_name,
+        valueInputOption="RAW",
+        body=body
+    ).execute()
+
+# Home Page
+if selected == "Home":
+    st.markdown(
+        """
+        ### Welcome to DataScribe
+        **DataScribe** is an AI-powered tool designed to extract structured information from the web 
+        based on entities in your data file. Start by uploading a CSV or Google Sheet and defining a 
+        custom search query.
+        """
+    )
+    st.image("https://via.placeholder.com/1200x400.png?text=DataScribe+AI+Agent+Dashboard")  # Placeholder banner image
+
 # Upload Data Section
-if selected == "Upload Data":
+elif selected == "Upload Data":
     st.header("Upload or Connect Your Data")
+
+    # CSV Upload
     data_source = st.radio("Choose data source:", ["CSV File", "Google Sheets"])
 
     if data_source == "CSV File":
@@ -89,14 +143,16 @@ if selected == "Upload Data":
 # Define Query Section
 elif selected == "Define Query":
     st.header("Define Your Custom Query")
+
     if st.session_state["data"] is not None:
         column_selection = st.selectbox("Select the primary column for entities", options=st.session_state["data"].columns)
         query_template = st.text_input("Define your query template", "Get me the email for {company}")
         st.session_state["query_template"] = query_template
-        st.session_state["column_selection"] = column_selection
+        st.session_state["column_selection"] = column_selection  # Store column selection in session state
 
         st.write("### Example query preview")
         if column_selection:
+            # Convert sample_entity to string to avoid replace errors
             sample_entity = str(st.session_state["data"][column_selection].iloc[0])
             example_query = query_template.replace("{company}", sample_entity)
             st.code(example_query)
@@ -106,31 +162,52 @@ elif selected == "Define Query":
 # Extract Information Section with Progress Bar
 elif selected == "Extract Information":
     st.header("Extract Information")
+
     if st.session_state.get("query_template") and st.session_state["data"] is not None and st.session_state["column_selection"] is not None:
         st.write("Data extraction is in progress. This may take a few moments.")
+
+        # Progress bar initialization
         progress_bar = st.progress(0)
         column_selection = st.session_state["column_selection"]
         progress_step = 1.0 / len(st.session_state["data"][column_selection])
 
         results = []
         for i, entity in enumerate(st.session_state["data"][column_selection]):
+            # Prepare the prompt for the model
             user_message = st.session_state["query_template"].replace("{company}", str(entity))
+            
+            # Get response and append to results
             result_text = get_llm_response(entity, user_message)
-            results.append({"Entity": entity, "Extracted Information": result_text})
+            results.append({"Entity": entity, "Extracted Information": result_text})  # Consistent key
+
+            # Update the progress bar
             progress_bar.progress((i + 1) * progress_step)
 
+        # Save and display results
         st.session_state["results"] = pd.DataFrame(results)
         st.write("### Extracted Information")
         st.dataframe(st.session_state["results"])
 
-# View & Download Section
+# View & Download Section with Google Sheets Update
 elif selected == "View & Download":
     st.header("View and Download Results")
+
     if st.session_state["results"] is not None:
         st.write("### Extracted Data Table")
         st.dataframe(st.session_state["results"])
 
+        # Download as CSV
         csv_data = st.session_state["results"].to_csv(index=False)
         st.download_button("Download as CSV", csv_data, "datascribe_results.csv", "text/csv")
+
+        # Option to update Google Sheet
+        sheet_id = st.text_input("Enter Google Sheet ID to update with results")
+        range_name = st.text_input("Enter range (e.g., Sheet1!A1)")
+        if st.button("Update Google Sheet"):
+            try:
+                update_google_sheet(sheet_id, range_name, st.session_state["results"])
+                st.success("Google Sheet updated successfully!")
+            except Exception as e:
+                st.error(f"Failed to update Google Sheet: {e}")
     else:
         st.warning("No data available to view or download.")