Gonalb commited on
Commit
660e128
·
1 Parent(s): d765e31
Files changed (1) hide show
  1. app.py +29 -159
app.py CHANGED
@@ -2,11 +2,6 @@ import chainlit as cl
2
  import pandas as pd
3
  import time
4
  from typing import Dict, Any
5
- import os
6
- import json
7
- from datetime import datetime
8
- import hashlib
9
- import uuid
10
 
11
  from agents.table_selection import table_selection_agent
12
  from agents.data_retrieval import sample_data_retrieval_agent
@@ -16,167 +11,42 @@ from agents.execution import execution_agent
16
  from utils.bigquery_utils import init_bigquery_connection
17
  from utils.feedback_utils import save_feedback_to_bigquery
18
 
19
- # File to store chat history
20
- HISTORY_DIR = "chat_history"
21
- os.makedirs(HISTORY_DIR, exist_ok=True)
22
-
23
- # Function to get or create a user ID
24
- def get_user_id():
25
- # Check if user already has an ID
26
- user_id = cl.user_session.get("user_id")
27
-
28
- if not user_id:
29
- # Generate a new user ID if none exists
30
- # In a real app, you might use authentication or cookies
31
- user_id = str(uuid.uuid4())
32
- cl.user_session.set("user_id", user_id)
33
-
34
- # Create a user directory for this user
35
- user_dir = os.path.join(HISTORY_DIR, user_id)
36
- os.makedirs(user_dir, exist_ok=True)
37
-
38
- return user_id
39
-
40
- # Function to save chat history
41
- def save_chat_history(user_id, conversation_id, user_message, assistant_response):
42
- user_dir = os.path.join(HISTORY_DIR, user_id)
43
- os.makedirs(user_dir, exist_ok=True)
44
-
45
- history_file = os.path.join(user_dir, f"{conversation_id}.json")
46
-
47
- # Load existing history if it exists
48
- if os.path.exists(history_file):
49
- with open(history_file, 'r') as f:
50
- history = json.load(f)
51
- else:
52
- history = {
53
- "conversation_id": conversation_id,
54
- "created_at": datetime.now().isoformat(),
55
- "messages": []
56
- }
57
-
58
- # Add new messages
59
- timestamp = datetime.now().isoformat()
60
- history["messages"].append({
61
- "timestamp": timestamp,
62
- "user": user_message,
63
- "assistant": assistant_response
64
- })
65
-
66
- # Save updated history
67
- with open(history_file, 'w') as f:
68
- json.dump(history, f, indent=2)
69
-
70
- # Function to load chat history for a user
71
- def load_chat_history(user_id, conversation_id=None):
72
- user_dir = os.path.join(HISTORY_DIR, user_id)
73
-
74
- if not os.path.exists(user_dir):
75
- return []
76
-
77
- if conversation_id:
78
- # Load specific conversation
79
- history_file = os.path.join(user_dir, f"{conversation_id}.json")
80
- if os.path.exists(history_file):
81
- with open(history_file, 'r') as f:
82
- return json.load(f)
83
- return None
84
- else:
85
- # Load all conversations (just metadata)
86
- conversations = []
87
- for filename in os.listdir(user_dir):
88
- if filename.endswith('.json'):
89
- with open(os.path.join(user_dir, filename), 'r') as f:
90
- data = json.load(f)
91
- conversations.append({
92
- "conversation_id": data["conversation_id"],
93
- "created_at": data["created_at"],
94
- "message_count": len(data["messages"])
95
- })
96
- return sorted(conversations, key=lambda x: x["created_at"], reverse=True)
97
-
98
  @cl.on_chat_start
99
- async def start():
100
- # Get or create user ID
101
- user_id = get_user_id()
102
-
103
- # Create a new conversation ID for this session
104
- conversation_id = str(uuid.uuid4())
105
- cl.user_session.set("conversation_id", conversation_id)
106
 
107
- # Load previous conversations for this user
108
- conversations = load_chat_history(user_id)
109
 
110
- if conversations:
111
- # Display a welcome back message
112
- await cl.Message(
113
- content=f"Welcome back! You have {len(conversations)} previous conversations.",
114
- ).send()
115
-
116
- # Create a dropdown to select previous conversations
117
- options = [{"value": conv["conversation_id"], "label": f"Conversation from {conv['created_at'][:10]} ({conv['message_count']} messages)"} for conv in conversations]
118
- options.insert(0, {"value": "new", "label": "Start a new conversation"})
119
-
120
- res = await cl.AskUserMessage(
121
- content="Would you like to continue a previous conversation or start a new one?",
122
- timeout=30,
123
- select={"options": options, "value": "new"}
124
- ).send()
125
-
126
- if res and res.get("value") != "new":
127
- # User selected a previous conversation
128
- selected_conv_id = res.get("value")
129
- cl.user_session.set("conversation_id", selected_conv_id)
130
-
131
- # Load the selected conversation
132
- conversation = load_chat_history(user_id, selected_conv_id)
133
- if conversation:
134
- # Display previous messages
135
- await cl.Message(content="Loading your previous conversation...").send()
136
- for msg in conversation["messages"]:
137
- await cl.Message(content=msg["user"], author="You").send()
138
- await cl.Message(content=msg["assistant"]).send()
139
- else:
140
- # First time user
141
- await cl.Message(
142
- content="Welcome to the E-commerce Analytics Assistant! How can I help you today?",
143
- ).send()
144
-
145
- # Add option to clear history
146
- await cl.Action(name="clear_history", label="Clear Chat History", description="Erase all previous conversation history").send()
147
-
148
- @cl.on_message
149
- async def on_message(message: cl.Message):
150
- # Get user and conversation IDs
151
- user_id = cl.user_session.get("user_id")
152
- conversation_id = cl.user_session.get("conversation_id")
153
-
154
- # Process the message and generate a response
155
- # Replace this with your actual processing logic
156
- response = f"You said: {message.content}"
157
-
158
- # Send the response
159
- await cl.Message(content=response).send()
160
 
161
- # Save to chat history
162
- save_chat_history(user_id, conversation_id, message.content, response)
163
-
164
- @cl.action_callback("clear_history")
165
- async def clear_history_callback(action):
166
- user_id = cl.user_session.get("user_id")
167
- conversation_id = cl.user_session.get("conversation_id")
168
 
169
- # Clear specific conversation or all conversations
170
- user_dir = os.path.join(HISTORY_DIR, user_id)
171
- history_file = os.path.join(user_dir, f"{conversation_id}.json")
 
 
 
172
 
173
- if os.path.exists(history_file):
174
- os.remove(history_file)
175
- await cl.Message(content="Your current conversation history has been cleared.").send()
176
 
177
- # Create a new conversation ID
178
- new_conversation_id = str(uuid.uuid4())
179
- cl.user_session.set("conversation_id", new_conversation_id)
 
180
 
181
  @cl.on_message
182
  async def on_message(message: cl.Message):
 
2
  import pandas as pd
3
  import time
4
  from typing import Dict, Any
 
 
 
 
 
5
 
6
  from agents.table_selection import table_selection_agent
7
  from agents.data_retrieval import sample_data_retrieval_agent
 
11
  from utils.bigquery_utils import init_bigquery_connection
12
  from utils.feedback_utils import save_feedback_to_bigquery
13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  @cl.on_chat_start
15
+ async def on_chat_start():
16
+ """Initialize the chat session."""
17
+ # Initialize BigQuery client
18
+ client = init_bigquery_connection()
 
 
 
19
 
20
+ # Store the client in the user session
21
+ cl.user_session.set("client", client)
22
 
23
+ # Send a welcome message
24
+ await cl.Message(
25
+ content="👋 Welcome to the Natural Language to SQL Query Assistant! Ask me any question about your e-commerce data.",
26
+ author="SQL Assistant"
27
+ ).send()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
28
 
29
+ # Add some example questions without using actions
30
+ await cl.Message(
31
+ content="Here are some example questions you can ask:",
32
+ author="SQL Assistant"
33
+ ).send()
 
 
34
 
35
+ examples = [
36
+ "What are the top 5 products by revenue?",
37
+ "How many orders were placed in the last month?",
38
+ "Which customers spent the most in 2023?",
39
+ "What is the average order value by product category?"
40
+ ]
41
 
42
+ # Display all examples in a single message
43
+ examples_text = "\n\n".join([f"• {example}" for example in examples])
44
+ examples_text += "\n\n(You can copy and paste any of these examples to try them out)"
45
 
46
+ await cl.Message(
47
+ content=examples_text,
48
+ author="SQL Assistant"
49
+ ).send()
50
 
51
  @cl.on_message
52
  async def on_message(message: cl.Message):