Update app.py

app.py

@@ -2,19 +2,13 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 import os
-import json
-from pathlib import Path
-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 (
     InfoSQLDatabaseTool,
     ListSQLDatabaseTool,
-    QuerySQLCheckerTool,
     QuerySQLDataBaseTool,
 )
 from langchain_community.utilities.sql_database import SQLDatabase
@@ -24,140 +18,143 @@ import tempfile
 st.title("Blah Blah App 🚀")
 st.write("Analyze datasets using natural language queries.")
 
-# Initialize LLM
-llm = None
+# LLM Initialization
+def initialize_llm(model_choice):
+    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.")
+            return None
+        return 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.")
+            return None
+        return ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
 
-# 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
-
-# Dataset Input
-input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
-if input_option == "Use Hugging Face Dataset":
-    dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
-    if st.button("Load Dataset"):
-        try:
-            with st.spinner("Loading dataset..."):
+llm = initialize_llm(model_choice)
+
+# Dataset Loading 
+def load_dataset_into_session():
+    input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
+    if input_option == "Use Hugging Face Dataset":
+        dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
+        if st.button("Load Dataset"):
+            try:
                 dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
                 st.session_state.df = pd.DataFrame(dataset)
                 st.success(f"Dataset '{dataset_name}' loaded successfully!")
                 st.dataframe(st.session_state.df.head())
-        except Exception as e:
-            st.error(f"Error: {e}")
-elif input_option == "Upload CSV File":
-    uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
-    if uploaded_file:
-        st.session_state.df = pd.read_csv(uploaded_file)
-        st.success("File uploaded successfully!")
-        st.dataframe(st.session_state.df.head())
+            except Exception as e:
+                st.error(f"Error: {e}")
+    elif input_option == "Upload CSV File":
+        uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
+        if uploaded_file:
+            st.session_state.df = pd.read_csv(uploaded_file)
+            st.success("File uploaded successfully!")
+            st.dataframe(st.session_state.df.head())
 
+if "df" not in st.session_state:
+    st.session_state.df = None
+load_dataset_into_session()
 
-if st.session_state.df is not None:
-    # Database setup
+# Database Initialization 
+def initialize_database(df):
     temp_dir = tempfile.TemporaryDirectory()
     db_path = os.path.join(temp_dir.name, "patent_data.db")
     connection = sqlite3.connect(db_path)
-    st.session_state.df.to_sql("patents", connection, if_exists="replace", index=False)
+    df.to_sql("patents", connection, if_exists="replace", index=False)
     db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
+    return db, temp_dir
 
-    # SQL Tools
+# SQL Tools
+def create_sql_tools(db):
     @tool("list_tables")
     def list_tables() -> str:
-        """List all tables in the patent database."""
         return ListSQLDatabaseTool(db=db).invoke("")
 
     @tool("tables_schema")
     def tables_schema(tables: str) -> str:
-        """Get schema and sample rows for given tables."""
         return InfoSQLDatabaseTool(db=db).invoke(tables)
 
     @tool("execute_sql")
     def execute_sql(sql_query: str) -> str:
-        """Execute a SQL query against the patent database."""
         return QuerySQLDataBaseTool(db=db).invoke(sql_query)
 
-    # --- CrewAI Agents for Patent Analysis ---
-    patent_sql_dev = Agent(
+    return list_tables, tables_schema, execute_sql
+
+# Agent Initialization
+def initialize_agents(llm, tools):
+    list_tables, tables_schema, execute_sql = tools
+
+    sql_agent = Agent(
         role="Patent Data Analyst",
         goal="Extract patent data using optimized SQL queries.",
-        backstory="An expert in writing optimized SQL queries for complex patent databases.",
+        backstory="Expert in optimized SQL for patent databases.",
         llm=llm,
         tools=[list_tables, tables_schema, execute_sql],
     )
 
-    patent_data_analyst = Agent(
+    analyst_agent = Agent(
         role="Patent Data Analyst",
         goal="Analyze the data and produce insights.",
-        backstory="A seasoned analyst who identifies trends and patterns in datasets.",
+        backstory="Data analyst identifying trends.",
         llm=llm,
     )
 
-    patent_report_writer = Agent(
+    writer_agent = Agent(
         role="Patent Report Writer",
-        goal="Summarize patent insights into a clear report.",
-        backstory="Expert in summarizing patent data insights into comprehensive reports.",
+        goal="Summarize patent insights into a report.",
+        backstory="Expert in clear, concise reporting.",
         llm=llm,
     )
 
-    # --- Crew Tasks ---
-    extract_data = Task(
+    return sql_agent, analyst_agent, writer_agent
+
+# Crew and Tasks Setup
+def setup_crew(sql_agent, analyst_agent, writer_agent):
+    extract_task = Task(
         description="Extract patents related to the query: {query}.",
         expected_output="Patent data matching the query.",
-        agent=patent_sql_dev,
+        agent=sql_agent,
     )
 
-    analyze_data = Task(
-        description="Analyze the extracted patent data for query: {query}.",
+    analyze_task = Task(
+        description="Analyze the extracted patent data.",
         expected_output="Analysis text summarizing findings.",
-        agent=patent_data_analyst,
-        context=[extract_data],
+        agent=analyst_agent,
+        context=[extract_task],
     )
 
-    write_report = Task(
-        description="Summarize analysis into an executive report.",
+    report_task = Task(
+        description="Summarize analysis into a report.",
         expected_output="Markdown report of insights.",
-        agent=patent_report_writer,
-        context=[analyze_data],
+        agent=writer_agent,
+        context=[analyze_task],
     )
 
-    # Assemble Crew
-    crew = Crew(
-        agents=[patent_sql_dev, patent_data_analyst, patent_report_writer],
-        tasks=[extract_data, analyze_data, write_report],
+    return Crew(
+        agents=[sql_agent, analyst_agent, writer_agent],
+        tasks=[extract_task, analyze_task, report_task],
         process=Process.sequential,
         verbose=True,
     )
 
-    #Query Input for Patent Analysis
+    # Execution Flow
+if st.session_state.df is not None:
+    db, temp_dir = initialize_database(st.session_state.df)
+    tools = create_sql_tools(db)
+    sql_agent, analyst_agent, writer_agent = initialize_agents(llm, tools)
+    crew = setup_crew(sql_agent, analyst_agent, writer_agent)
+
     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..."):
-            inputs = {"query": query}
-            result = crew.kickoff(inputs=inputs)
+            result = crew.kickoff(inputs={"query": query})
             st.markdown("### 📊 Patent Analysis Report")
             st.markdown(result)
-
-    temp_dir.cleanup()
 else:
-    st.info("Please load a patent dataset to proceed.")
+    st.info("Please load a patent dataset to proceed.")