app.py

import streamlit as st
from streamlit_option_menu import option_menu
import pandas as pd
import os
from google.oauth2 import service_account
from googleapiclient.discovery import build
from streamlit_chat import message as st_message
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 re

warnings.filterwarnings("ignore", category=DeprecationWarning)

# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
os.environ['GROQ_API_KEY'] = GROQ_API_KEY
llm = ChatGroq(model="llama-3.1-70b-versatile")

PROMPT_TEMPLATE = """
You are an expert information extraction assistant designed to obtain specific details from the web and external sources.
You’ll be provided with an entity name and a query that specifies the type of information needed about that entity.
Please follow the instructions carefully and return only the most relevant, accurate information.

#### Entity Name: {entity}
#### Query: {query}

Instructions:
1. Extract the information directly related to the entity.
2. If available, include only verified, publicly accessible data.
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.
"""

# Set up the page
st.set_page_config(page_title="DataScribe", page_icon=":notebook_with_decorative_cover:", layout="wide")

# Sidebar navigation
with st.sidebar:
    selected = option_menu(
        "DataScribe Menu",
        ["Home", "Upload Data", "Define Query", "Extract Information", "View & Download"],
        icons=["house", "cloud-upload", "gear", "search", "table"],
        menu_icon="cast",
        default_index=0
    )

# Main header
st.title("DataScribe: AI-Powered Information Extractor")

# Initialize session states for data and results
if "data" not in st.session_state:
    st.session_state["data"] = None
if "results" not in st.session_state:
    st.session_state["results"] = None
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
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":
        uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])
        if uploaded_file:
            st.session_state["data"] = pd.read_csv(uploaded_file)
            st.write("### Preview of Uploaded Data")
            st.dataframe(st.session_state["data"].head())

    elif data_source == "Google Sheets":
        sheet_id = st.text_input("Enter Google Sheet ID")
        range_name = st.text_input("Enter the data range (e.g., Sheet1!A1:C100)")
        if st.button("Fetch Data"):
            if sheet_id and range_name:
                st.session_state["data"] = get_google_sheet_data(sheet_id, range_name)
                st.write("### Preview of Google Sheets Data")
                st.dataframe(st.session_state["data"].head())
            else:
                st.warning("Please enter both the Google Sheet ID and range.")

# 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  # 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)
    else:
        st.warning("Please upload data first.")

# 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))
            formatted_prompt = PROMPT_TEMPLATE.format(entity=entity, query=user_message)
            
            # Append user message to the flow history
            st.session_state.flowmessages.append(HumanMessage(content=user_message))

            # Generate response from the model
            response = llm([SystemMessage(content=formatted_prompt)])

            # Collect the model's response
            result_text = response[0].content if response else "Information not available"
            results.append({"Entity": entity, "Extracted Information": result_text})

            # 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 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.")