DrishtiSharma commited on
Commit
8ddb996
Β·
verified Β·
1 Parent(s): ff795fd

Create flawed_crew.py

Browse files
Files changed (1) hide show
  1. mylab/flawed_crew.py +374 -0
mylab/flawed_crew.py ADDED
@@ -0,0 +1,374 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import pandas as pd
3
+ import sqlite3
4
+ import os
5
+ import io
6
+ import json
7
+ from pathlib import Path
8
+ import tempfile
9
+ from fpdf import FPDF
10
+ import plotly.express as px
11
+ from datetime import datetime, timezone
12
+ from crewai import Agent, Crew, Process, Task
13
+ from crewai.tools import tool
14
+ from langchain_groq import ChatGroq
15
+ from langchain_openai import ChatOpenAI
16
+ from langchain.schema.output import LLMResult
17
+ from langchain_community.tools.sql_database.tool import (
18
+ InfoSQLDatabaseTool,
19
+ ListSQLDatabaseTool,
20
+ QuerySQLCheckerTool,
21
+ QuerySQLDataBaseTool,
22
+ )
23
+ from langchain_community.utilities.sql_database import SQLDatabase
24
+ from datasets import load_dataset
25
+ import tempfile
26
+
27
+ st.title("SQL-RAG Using CrewAI πŸš€")
28
+ st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
29
+
30
+ # Initialize LLM
31
+ llm = None
32
+
33
+ # Model Selection
34
+ model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
35
+
36
+ # API Key Validation and LLM Initialization
37
+ groq_api_key = os.getenv("GROQ_API_KEY")
38
+ openai_api_key = os.getenv("OPENAI_API_KEY")
39
+
40
+ if model_choice == "llama-3.3-70b":
41
+ if not groq_api_key:
42
+ st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
43
+ llm = None
44
+ else:
45
+ llm = ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
46
+ elif model_choice == "GPT-4o":
47
+ if not openai_api_key:
48
+ st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
49
+ llm = None
50
+ else:
51
+ llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
52
+
53
+ # Initialize session state for data persistence
54
+ if "df" not in st.session_state:
55
+ st.session_state.df = None
56
+ if "show_preview" not in st.session_state:
57
+ st.session_state.show_preview = False
58
+
59
+ # Dataset Input
60
+ input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
61
+
62
+ if input_option == "Use Hugging Face Dataset":
63
+ dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="Einstellung/demo-salaries")
64
+ if st.button("Load Dataset"):
65
+ try:
66
+ with st.spinner("Loading dataset..."):
67
+ dataset = load_dataset(dataset_name, split="train")
68
+ st.session_state.df = pd.DataFrame(dataset)
69
+ st.session_state.show_preview = True # Show preview after loading
70
+ st.success(f"Dataset '{dataset_name}' loaded successfully!")
71
+ except Exception as e:
72
+ st.error(f"Error: {e}")
73
+
74
+ elif input_option == "Upload CSV File":
75
+ uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
76
+ if uploaded_file:
77
+ try:
78
+ st.session_state.df = pd.read_csv(uploaded_file)
79
+ st.session_state.show_preview = True # Show preview after loading
80
+ st.success("File uploaded successfully!")
81
+ except Exception as e:
82
+ st.error(f"Error loading file: {e}")
83
+
84
+ # Show Dataset Preview Only After Loading
85
+ if st.session_state.df is not None and st.session_state.show_preview:
86
+ st.subheader("πŸ“‚ Dataset Preview")
87
+ st.dataframe(st.session_state.df.head())
88
+
89
+ # Helper Function to Create a PDF Report with Visualizations and Descriptions
90
+ def create_pdf_report_with_viz(report, conclusion, visualizations):
91
+ pdf = FPDF()
92
+ pdf.set_auto_page_break(auto=True, margin=15)
93
+ pdf.add_page()
94
+ pdf.set_font("Arial", size=12)
95
+
96
+ # Title
97
+ pdf.set_font("Arial", style="B", size=18)
98
+ pdf.cell(0, 10, "πŸ“Š Analysis Report", ln=True, align="C")
99
+ pdf.ln(10)
100
+
101
+ # Report Content
102
+ pdf.set_font("Arial", style="B", size=14)
103
+ pdf.cell(0, 10, "Analysis", ln=True)
104
+ pdf.set_font("Arial", size=12)
105
+ pdf.multi_cell(0, 10, report)
106
+
107
+ pdf.ln(10)
108
+ pdf.set_font("Arial", style="B", size=14)
109
+ pdf.cell(0, 10, "Conclusion", ln=True)
110
+ pdf.set_font("Arial", size=12)
111
+ pdf.multi_cell(0, 10, conclusion)
112
+
113
+ # Add Visualizations
114
+ pdf.add_page()
115
+ pdf.set_font("Arial", style="B", size=16)
116
+ pdf.cell(0, 10, "πŸ“ˆ Visualizations", ln=True)
117
+ pdf.ln(5)
118
+
119
+ with tempfile.TemporaryDirectory() as temp_dir:
120
+ for i, fig in enumerate(visualizations, start=1):
121
+ fig_title = fig.layout.title.text if fig.layout.title.text else f"Visualization {i}"
122
+ x_axis = fig.layout.xaxis.title.text if fig.layout.xaxis.title.text else "X-axis"
123
+ y_axis = fig.layout.yaxis.title.text if fig.layout.yaxis.title.text else "Y-axis"
124
+
125
+ # Save each visualization as a PNG image
126
+ img_path = os.path.join(temp_dir, f"viz_{i}.png")
127
+ fig.write_image(img_path)
128
+
129
+ # Insert Title and Description
130
+ pdf.set_font("Arial", style="B", size=14)
131
+ pdf.multi_cell(0, 10, f"{i}. {fig_title}")
132
+ pdf.set_font("Arial", size=12)
133
+ pdf.multi_cell(0, 10, f"X-axis: {x_axis} | Y-axis: {y_axis}")
134
+ pdf.ln(3)
135
+
136
+ # Embed Visualization
137
+ pdf.image(img_path, w=170)
138
+ pdf.ln(10)
139
+
140
+ # Save PDF
141
+ temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
142
+ pdf.output(temp_pdf.name)
143
+
144
+ return temp_pdf
145
+
146
+
147
+ # Helper function to create a plain text report with visualization summaries
148
+ def create_text_report_with_viz(report, conclusion, visualizations):
149
+ content = f"### Analysis Report\n\n{report}\n\n### Conclusion\n\n{conclusion}\n\n### Visualizations\n"
150
+
151
+ # Dynamically add descriptions for each visualization
152
+ for i, fig in enumerate(visualizations, start=1):
153
+ # Extract the title of the Plotly figure if available
154
+ fig_title = fig.layout.title.text if fig.layout.title.text else f"Visualization {i}"
155
+
156
+ # Add title and figure details
157
+ content += f"\n{i}. {fig_title}\n"
158
+
159
+ # Extract x and y axis titles if available
160
+ x_axis = fig.layout.xaxis.title.text if fig.layout.xaxis.title.text else "X-axis"
161
+ y_axis = fig.layout.yaxis.title.text if fig.layout.yaxis.title.text else "Y-axis"
162
+
163
+ content += f" - X-axis: {x_axis}\n"
164
+ content += f" - Y-axis: {y_axis}\n"
165
+
166
+ # If figure has data, summarize it
167
+ if fig.data:
168
+ trace_types = set(trace.type for trace in fig.data)
169
+ content += f" - Chart Type(s): {', '.join(trace_types)}\n"
170
+ else:
171
+ content += " - No data available in this visualization.\n"
172
+
173
+ # Return the content as a downloadable text stream
174
+ return io.BytesIO(content.encode("utf-8"))
175
+
176
+
177
+ # SQL-RAG Analysis
178
+ if st.session_state.df is not None:
179
+ temp_dir = tempfile.TemporaryDirectory()
180
+ db_path = os.path.join(temp_dir.name, "data.db")
181
+ connection = sqlite3.connect(db_path)
182
+ st.session_state.df.to_sql("salaries", connection, if_exists="replace", index=False)
183
+ db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
184
+
185
+ @tool("list_tables")
186
+ def list_tables() -> str:
187
+ """List all tables in the database."""
188
+ return ListSQLDatabaseTool(db=db).invoke("")
189
+
190
+ @tool("tables_schema")
191
+ def tables_schema(tables: str) -> str:
192
+ """Get the schema and sample rows for the specified tables."""
193
+ return InfoSQLDatabaseTool(db=db).invoke(tables)
194
+
195
+ @tool("execute_sql")
196
+ def execute_sql(sql_query: str) -> str:
197
+ """Execute a SQL query against the database and return the results."""
198
+ return QuerySQLDataBaseTool(db=db).invoke(sql_query)
199
+
200
+ @tool("check_sql")
201
+ def check_sql(sql_query: str) -> str:
202
+ """Validate the SQL query syntax and structure before execution."""
203
+ return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
204
+
205
+ # Agents for SQL data extraction and analysis
206
+ sql_dev = Agent(
207
+ role="Senior Database Developer",
208
+ goal="Extract data using optimized SQL queries.",
209
+ backstory="An expert in writing optimized SQL queries for complex databases.",
210
+ llm=llm,
211
+ tools=[list_tables, tables_schema, execute_sql, check_sql],
212
+ )
213
+
214
+ data_analyst = Agent(
215
+ role="Senior Data Analyst",
216
+ goal="Analyze the data and produce insights.",
217
+ backstory="A seasoned analyst who identifies trends and patterns in datasets.",
218
+ llm=llm,
219
+ )
220
+
221
+ report_writer = Agent(
222
+ role="Technical Report Writer",
223
+ goal="Write a structured report with Introduction and Key Insights. DO NOT include any Conclusion or Summary.",
224
+ backstory="Specializes in detailed analytical reports without conclusions.",
225
+ llm=llm,
226
+ )
227
+
228
+ conclusion_writer = Agent(
229
+ role="Conclusion Specialist",
230
+ goal="Summarize findings into a clear and concise 3-5 line Conclusion highlighting only the most important insights.",
231
+ backstory="An expert in crafting impactful and clear conclusions.",
232
+ llm=llm,
233
+ )
234
+
235
+ # Define tasks for report and conclusion
236
+ extract_data = Task(
237
+ description="Extract data based on the query: {query}.",
238
+ expected_output="Database results matching the query.",
239
+ agent=sql_dev,
240
+ )
241
+
242
+ analyze_data = Task(
243
+ description="Analyze the extracted data for query: {query}.",
244
+ expected_output="Key Insights and Analysis without any Introduction or Conclusion.",
245
+ agent=data_analyst,
246
+ context=[extract_data],
247
+ )
248
+
249
+ write_report = Task(
250
+ description="Write the analysis report with Introduction and Key Insights. DO NOT include any Conclusion or Summary.",
251
+ expected_output="Markdown-formatted report excluding Conclusion.",
252
+ agent=report_writer,
253
+ context=[analyze_data],
254
+ )
255
+
256
+ write_conclusion = Task(
257
+ description="Summarize the key findings in 3-5 impactful lines, highlighting the maximum, minimum, and average salaries."
258
+ "Emphasize significant insights on salary distribution and influential compensation trends for strategic decision-making.",
259
+ expected_output="Markdown-formatted Conclusion section with key insights and statistics.",
260
+ agent=conclusion_writer,
261
+ context=[analyze_data],
262
+ )
263
+
264
+ # Separate Crews for report and conclusion
265
+ crew_report = Crew(
266
+ agents=[sql_dev, data_analyst, report_writer],
267
+ tasks=[extract_data, analyze_data, write_report],
268
+ process=Process.sequential,
269
+ verbose=True,
270
+ )
271
+
272
+ crew_conclusion = Crew(
273
+ agents=[data_analyst, conclusion_writer],
274
+ tasks=[write_conclusion],
275
+ process=Process.sequential,
276
+ verbose=True,
277
+ )
278
+
279
+ # Tabs for Query Results and Visualizations
280
+ tab1, tab2 = st.tabs(["πŸ” Query Insights + Viz", "πŸ“Š Full Data Viz"])
281
+
282
+ # Query Insights + Visualization
283
+ with tab1:
284
+ query = st.text_area("Enter Query:", value="Provide insights into the salary of a Principal Data Scientist.")
285
+ if st.button("Submit Query"):
286
+ with st.spinner("Processing query..."):
287
+ # Step 1: Generate the analysis report
288
+ report_inputs = {"query": query + " Provide detailed analysis but DO NOT include Conclusion."}
289
+ report_result = crew_report.kickoff(inputs=report_inputs)
290
+
291
+ # Step 2: Generate only the concise conclusion
292
+ conclusion_inputs = {"query": query + " Provide ONLY the most important insights in 3-5 concise lines."}
293
+ conclusion_result = crew_conclusion.kickoff(inputs=conclusion_inputs)
294
+
295
+ # Step 3: Display the report
296
+ #st.markdown("### Analysis Report:")
297
+ st.markdown(report_result if report_result else "⚠️ No Report Generated.")
298
+
299
+ # Step 4: Generate Visualizations
300
+ visualizations = []
301
+
302
+ fig_salary = px.box(st.session_state.df, x="job_title", y="salary_in_usd",
303
+ title="Salary Distribution by Job Title")
304
+ visualizations.append(fig_salary)
305
+
306
+ fig_experience = px.bar(
307
+ st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
308
+ x="experience_level", y="salary_in_usd",
309
+ title="Average Salary by Experience Level"
310
+ )
311
+ visualizations.append(fig_experience)
312
+
313
+ fig_employment = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
314
+ title="Salary Distribution by Employment Type")
315
+ visualizations.append(fig_employment)
316
+
317
+ # Step 5: Insert Visual Insights
318
+ st.markdown("### Visual Insights")
319
+ for fig in visualizations:
320
+ st.plotly_chart(fig, use_container_width=True)
321
+
322
+ # Step 6: Display Concise Conclusion
323
+ #st.markdown("#### 6. Conclusion")
324
+ st.markdown(conclusion_result if conclusion_result else "⚠️ No Conclusion Generated.")
325
+
326
+ # Step 7: PDF and TXT Download Buttons for Query Insights + Viz
327
+ if report_result and conclusion_result and visualizations:
328
+ # PDF Download
329
+ pdf_file = create_pdf_report_with_viz(report_result, conclusion_result, visualizations)
330
+ with open(pdf_file.name, "rb") as f:
331
+ st.download_button(
332
+ label="πŸ“₯ Download Full Report (PDF)",
333
+ data=f,
334
+ file_name="query_insights_report.pdf",
335
+ mime="application/pdf"
336
+ )
337
+
338
+ # TXT Download
339
+ text_file = create_text_report_with_viz(report_result, conclusion_result, visualizations)
340
+ st.download_button(
341
+ label="πŸ“₯ Download Full Report (TXT)",
342
+ data=text_file,
343
+ file_name="query_insights_report.txt",
344
+ mime="text/plain"
345
+ )
346
+
347
+
348
+ # Full Data Visualization Tab
349
+ with tab2:
350
+ st.subheader("πŸ“Š Comprehensive Data Visualizations")
351
+
352
+ fig1 = px.histogram(st.session_state.df, x="job_title", title="Job Title Frequency")
353
+ st.plotly_chart(fig1)
354
+
355
+ fig2 = px.bar(
356
+ st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
357
+ x="experience_level", y="salary_in_usd",
358
+ title="Average Salary by Experience Level"
359
+ )
360
+ st.plotly_chart(fig2)
361
+
362
+ fig3 = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
363
+ title="Salary Distribution by Employment Type")
364
+ st.plotly_chart(fig3)
365
+
366
+ temp_dir.cleanup()
367
+ else:
368
+ st.info("Please load a dataset to proceed.")
369
+
370
+
371
+ # Sidebar Reference
372
+ with st.sidebar:
373
+ st.header("πŸ“š Reference:")
374
+ 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)")