DrishtiSharma commited on
Commit
72ea0b4
Β·
verified Β·
1 Parent(s): 7832c6d

Create qt.py

Browse files
Files changed (1) hide show
  1. qt.py +236 -0
qt.py ADDED
@@ -0,0 +1,236 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import pandas as pd
3
+ import sqlite3
4
+ import os
5
+ import json
6
+ from pathlib import Path
7
+ import plotly.express as px
8
+ from datetime import datetime, timezone
9
+ from crewai import Agent, Crew, Process, Task
10
+ from crewai.tools import tool
11
+ from langchain_groq import ChatGroq
12
+ from langchain_openai import ChatOpenAI
13
+ from langchain.schema.output import LLMResult
14
+ from langchain_community.tools.sql_database.tool import (
15
+ InfoSQLDatabaseTool,
16
+ ListSQLDatabaseTool,
17
+ QuerySQLCheckerTool,
18
+ QuerySQLDataBaseTool,
19
+ )
20
+ from langchain_community.utilities.sql_database import SQLDatabase
21
+ from datasets import load_dataset
22
+ import tempfile
23
+
24
+ st.title("SQL-RAG Using CrewAI πŸš€")
25
+ st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
26
+
27
+ # Initialize LLM
28
+ llm = None
29
+
30
+ # Model Selection
31
+ model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
32
+
33
+ # API Key Validation and LLM Initialization
34
+ groq_api_key = os.getenv("GROQ_API_KEY")
35
+ openai_api_key = os.getenv("OPENAI_API_KEY")
36
+
37
+ if model_choice == "llama-3.3-70b":
38
+ if not groq_api_key:
39
+ st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
40
+ llm = None
41
+ else:
42
+ llm = ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
43
+ elif model_choice == "GPT-4o":
44
+ if not openai_api_key:
45
+ st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
46
+ llm = None
47
+ else:
48
+ llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
49
+
50
+ # Initialize session state for data persistence
51
+ if "df" not in st.session_state:
52
+ st.session_state.df = None
53
+ if "show_preview" not in st.session_state:
54
+ st.session_state.show_preview = False
55
+
56
+ # Dataset Input
57
+ input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
58
+
59
+ if input_option == "Use Hugging Face Dataset":
60
+ dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="Einstellung/demo-salaries")
61
+ if st.button("Load Dataset"):
62
+ try:
63
+ with st.spinner("Loading dataset..."):
64
+ dataset = load_dataset(dataset_name, split="train")
65
+ st.session_state.df = pd.DataFrame(dataset)
66
+ st.session_state.show_preview = True # Show preview after loading
67
+ st.success(f"Dataset '{dataset_name}' loaded successfully!")
68
+ except Exception as e:
69
+ st.error(f"Error: {e}")
70
+
71
+ elif input_option == "Upload CSV File":
72
+ uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
73
+ if uploaded_file:
74
+ try:
75
+ st.session_state.df = pd.read_csv(uploaded_file)
76
+ st.session_state.show_preview = True # Show preview after loading
77
+ st.success("File uploaded successfully!")
78
+ except Exception as e:
79
+ st.error(f"Error loading file: {e}")
80
+
81
+ # Show Dataset Preview Only After Loading
82
+ if st.session_state.df is not None and st.session_state.show_preview:
83
+ st.subheader("πŸ“‚ Dataset Preview")
84
+ st.dataframe(st.session_state.df.head())
85
+
86
+ # SQL-RAG Analysis
87
+ if st.session_state.df is not None:
88
+ temp_dir = tempfile.TemporaryDirectory()
89
+ db_path = os.path.join(temp_dir.name, "data.db")
90
+ connection = sqlite3.connect(db_path)
91
+ st.session_state.df.to_sql("salaries", connection, if_exists="replace", index=False)
92
+ db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
93
+
94
+ @tool("list_tables")
95
+ def list_tables() -> str:
96
+ """List all tables in the database."""
97
+ return ListSQLDatabaseTool(db=db).invoke("")
98
+
99
+ @tool("tables_schema")
100
+ def tables_schema(tables: str) -> str:
101
+ """Get the schema and sample rows for the specified tables."""
102
+ return InfoSQLDatabaseTool(db=db).invoke(tables)
103
+
104
+ @tool("execute_sql")
105
+ def execute_sql(sql_query: str) -> str:
106
+ """Execute a SQL query against the database and return the results."""
107
+ return QuerySQLDataBaseTool(db=db).invoke(sql_query)
108
+
109
+ @tool("check_sql")
110
+ def check_sql(sql_query: str) -> str:
111
+ """Validate the SQL query syntax and structure before execution."""
112
+ return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
113
+
114
+ sql_dev = Agent(
115
+ role="Senior Database Developer",
116
+ goal="Extract data using optimized SQL queries.",
117
+ backstory="An expert in writing optimized SQL queries for complex databases.",
118
+ llm=llm,
119
+ tools=[list_tables, tables_schema, execute_sql, check_sql],
120
+ )
121
+
122
+ data_analyst = Agent(
123
+ role="Senior Data Analyst",
124
+ goal="Analyze the data and produce insights.",
125
+ backstory="A seasoned analyst who identifies trends and patterns in datasets.",
126
+ llm=llm,
127
+ )
128
+
129
+ report_writer = Agent(
130
+ role="Technical Report Writer",
131
+ goal="Summarize the insights into a clear report.",
132
+ backstory="An expert in summarizing data insights into readable reports.",
133
+ llm=llm,
134
+ )
135
+
136
+ extract_data = Task(
137
+ description="Extract data based on the query: {query}.",
138
+ expected_output="Database results matching the query.",
139
+ agent=sql_dev,
140
+ )
141
+
142
+ analyze_data = Task(
143
+ description="Analyze the extracted data for query: {query}.",
144
+ expected_output="Analysis text summarizing findings (without a Conclusion section).",
145
+ agent=data_analyst,
146
+ context=[extract_data],
147
+ )
148
+
149
+ write_report = Task(
150
+ description="Summarize the analysis into an executive report without a Conclusion.",
151
+ expected_output="Markdown report of insights without Conclusion.",
152
+ agent=report_writer,
153
+ context=[analyze_data],
154
+ )
155
+
156
+ crew = Crew(
157
+ agents=[sql_dev, data_analyst, report_writer],
158
+ tasks=[extract_data, analyze_data, write_report],
159
+ process=Process.sequential,
160
+ verbose=True,
161
+ )
162
+
163
+ # Tabs for Query Results and General Insights
164
+ tab1, tab2 = st.tabs(["πŸ” Query Insights + Viz", "πŸ“Š Full Data Viz"])
165
+
166
+ # Tab 1: Query-Insights + Visualization
167
+ with tab1:
168
+ query = st.text_area("Enter Query:", value="Provide insights into the salary of a Principal Data Scientist.")
169
+ if st.button("Submit Query"):
170
+ with st.spinner("Processing query..."):
171
+ # Step 1: Generate Report without Conclusion
172
+ inputs = {"query": query + " Provide a detailed analysis but DO NOT include a Conclusion."}
173
+ report_result = crew.kickoff(inputs=inputs)
174
+
175
+ # Step 2: Generate only the Conclusion
176
+ conclusion_inputs = {"query": query + " Now, provide only the Conclusion for this analysis."}
177
+ conclusion_result = crew.kickoff(inputs=conclusion_inputs)
178
+
179
+ st.markdown("### Analysis Report:")
180
+
181
+ # Step 3: Generate relevant visualizations
182
+ visualizations = []
183
+
184
+ fig_salary = px.box(st.session_state.df, x="job_title", y="salary_in_usd",
185
+ title="Salary Distribution by Job Title")
186
+ visualizations.append(fig_salary)
187
+
188
+ fig_experience = px.bar(
189
+ st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
190
+ x="experience_level", y="salary_in_usd",
191
+ title="Average Salary by Experience Level"
192
+ )
193
+ visualizations.append(fig_experience)
194
+
195
+ fig_employment = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
196
+ title="Salary Distribution by Employment Type")
197
+ visualizations.append(fig_employment)
198
+
199
+ # Step 4: Display report without conclusion
200
+ st.markdown(report_result)
201
+
202
+ # Step 5: Insert Visual Insights
203
+ st.markdown("## πŸ“Š Visual Insights")
204
+ for fig in visualizations:
205
+ st.plotly_chart(fig, use_container_width=True)
206
+
207
+ # Step 6: Append the Conclusion
208
+ st.markdown("## Conclusion")
209
+ st.markdown(conclusion_result)
210
+
211
+ # Tab 2: Full Data Visualization
212
+ with tab2:
213
+ st.subheader("πŸ“Š Comprehensive Data Visualizations")
214
+
215
+ fig1 = px.histogram(st.session_state.df, x="job_title", title="Job Title Frequency")
216
+ st.plotly_chart(fig1)
217
+
218
+ fig2 = px.bar(
219
+ st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
220
+ x="experience_level", y="salary_in_usd",
221
+ title="Average Salary by Experience Level"
222
+ )
223
+ st.plotly_chart(fig2)
224
+
225
+ fig3 = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
226
+ title="Salary Distribution by Employment Type")
227
+ st.plotly_chart(fig3)
228
+
229
+ temp_dir.cleanup()
230
+ else:
231
+ st.info("Please load a dataset to proceed.")
232
+
233
+ # Sidebar Reference
234
+ with st.sidebar:
235
+ st.header("πŸ“š Reference:")
236
+ st.markdown("[SQL Agents w CrewAI & Llama 3 - Plaban Nayak](https://github.com/plaban1981/Agents/blob/main/SQL_Agents_with_CrewAI_and_Llama_3.ipynb)")