Divax
commited on
Commit
·
5f1c192
1
Parent(s):
d677138
files added
Browse files- app.py +205 -0
- requirements.txt +3 -0
app.py
ADDED
@@ -0,0 +1,205 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from groq import Groq
|
3 |
+
import sys
|
4 |
+
import io
|
5 |
+
import re
|
6 |
+
import traceback
|
7 |
+
from dotenv import load_dotenv
|
8 |
+
import os
|
9 |
+
|
10 |
+
load_dotenv()
|
11 |
+
|
12 |
+
# Set page configuration
|
13 |
+
st.set_page_config(
|
14 |
+
page_title="AI Assistant with Code Execution",
|
15 |
+
page_icon="🤖",
|
16 |
+
layout="wide",
|
17 |
+
)
|
18 |
+
|
19 |
+
def main():
|
20 |
+
st.title("AI Assistant with Code Execution")
|
21 |
+
st.write("Interact with an AI assistant that can execute code when needed.")
|
22 |
+
st.markdown("[Concept by cfahlgren1, check out his amazing work here](https://huggingface.co/spaces/cfahlgren1/qwen-2.5-code-interpreter)")
|
23 |
+
st.markdown("Thanks Groq for the super fast model and thanks cfahlgren1 for the prompt and this idea!")
|
24 |
+
|
25 |
+
# Sidebar settings
|
26 |
+
st.sidebar.title("Settings")
|
27 |
+
use_code_interpreter = st.sidebar.checkbox("Enable Code Interpreter", value=True)
|
28 |
+
reset_button = st.sidebar.button("Reset Chat")
|
29 |
+
|
30 |
+
# Initialize session state
|
31 |
+
if 'messages' not in st.session_state or reset_button:
|
32 |
+
st.session_state['messages'] = [
|
33 |
+
{
|
34 |
+
"role": "system",
|
35 |
+
"content": (
|
36 |
+
"The user will ask you a tricky question, your job is to write Python code to answer the question. \n\n" +
|
37 |
+
"Really think step by step before writing any code to ensure you're answering the question correctly. \n\n" +
|
38 |
+
"Respond with a markdown code block starting with ```python and ``` at the end. Make sure the code can be executed without any changes"
|
39 |
+
),
|
40 |
+
}
|
41 |
+
]
|
42 |
+
if reset_button:
|
43 |
+
st.session_state['user_input'] = ''
|
44 |
+
st.rerun()
|
45 |
+
|
46 |
+
# Display chat history (only user and assistant messages)
|
47 |
+
for message in st.session_state['messages']:
|
48 |
+
if message["role"] == "user":
|
49 |
+
with st.chat_message("user"):
|
50 |
+
st.markdown(message["content"])
|
51 |
+
elif message["role"] == "assistant":
|
52 |
+
with st.chat_message("assistant"):
|
53 |
+
st.markdown(message["content"])
|
54 |
+
# Internal messages (e.g., code execution results) are not displayed
|
55 |
+
|
56 |
+
# User input
|
57 |
+
if 'user_input' not in st.session_state:
|
58 |
+
st.session_state['user_input'] = ''
|
59 |
+
|
60 |
+
st.session_state['user_input'] = st.chat_input("Type your message")
|
61 |
+
if st.session_state['user_input']:
|
62 |
+
# Display user's message
|
63 |
+
with st.chat_message("user"):
|
64 |
+
st.markdown(st.session_state['user_input'])
|
65 |
+
# Add user's message to conversation history
|
66 |
+
st.session_state['messages'].append({"role": "user", "content": st.session_state['user_input']})
|
67 |
+
|
68 |
+
# Generate assistant's response
|
69 |
+
assistant_reply = get_assistant_response(st.session_state['messages'], use_code_interpreter)
|
70 |
+
|
71 |
+
# Display assistant's final answer
|
72 |
+
with st.chat_message("assistant"):
|
73 |
+
st.markdown(assistant_reply)
|
74 |
+
|
75 |
+
# Add assistant's final answer to conversation history
|
76 |
+
st.session_state['messages'].append({"role": "assistant", "content": assistant_reply})
|
77 |
+
|
78 |
+
# Clear user input
|
79 |
+
st.session_state['user_input'] = ''
|
80 |
+
|
81 |
+
def get_assistant_response(conversation, use_code_interpreter):
|
82 |
+
# Copy conversation to avoid modifying original
|
83 |
+
messages = conversation.copy()
|
84 |
+
|
85 |
+
# Initialize Groq client (Replace 'your_api_key' with your actual API key)
|
86 |
+
client = Groq()
|
87 |
+
|
88 |
+
# Get assistant's initial response (which may include code)
|
89 |
+
completion = client.chat.completions.create(
|
90 |
+
model="llama-3.2-3b-preview",
|
91 |
+
messages=messages,
|
92 |
+
temperature=1,
|
93 |
+
max_tokens=1024,
|
94 |
+
top_p=1,
|
95 |
+
stream=False,
|
96 |
+
stop=None,
|
97 |
+
)
|
98 |
+
assistant_reply = completion.choices[0].message.content
|
99 |
+
# print(assistant_reply)
|
100 |
+
# If code interpreter is enabled, check for code in assistant's reply
|
101 |
+
if use_code_interpreter:
|
102 |
+
code = extract_code(assistant_reply)
|
103 |
+
if code:
|
104 |
+
# Remove code from assistant's reply to avoid displaying it
|
105 |
+
assistant_reply_no_code = remove_code_blocks(assistant_reply)
|
106 |
+
|
107 |
+
# Execute the code
|
108 |
+
execution_result = execute_code_safely(code)
|
109 |
+
|
110 |
+
# Prepare internal message with execution result (not shown in UI chat)
|
111 |
+
user_result_prompt = (
|
112 |
+
f"The code you provided was executed and returned the result:\n{execution_result}\n"
|
113 |
+
"Use this result to provide a final answer to the user's question. "
|
114 |
+
"Do not mention the code or that you executed code."
|
115 |
+
)
|
116 |
+
|
117 |
+
# Add internal message to conversation history
|
118 |
+
messages.append({"role": "user", "content": user_result_prompt})
|
119 |
+
|
120 |
+
# Get assistant's final response using the execution result
|
121 |
+
completion = client.chat.completions.create(
|
122 |
+
model="llama-3.2-3b-preview",
|
123 |
+
messages=messages,
|
124 |
+
temperature=1,
|
125 |
+
max_tokens=1024,
|
126 |
+
top_p=1,
|
127 |
+
stream=False,
|
128 |
+
stop=None,
|
129 |
+
)
|
130 |
+
final_reply = completion.choices[0].message.content
|
131 |
+
# Return assistant's final answer
|
132 |
+
return final_reply.strip()
|
133 |
+
else:
|
134 |
+
# No code detected; return assistant's initial reply
|
135 |
+
return assistant_reply.strip()
|
136 |
+
else:
|
137 |
+
# Code interpreter disabled; return assistant's initial reply
|
138 |
+
return assistant_reply.strip()
|
139 |
+
|
140 |
+
def extract_code(text):
|
141 |
+
# Extract code from code blocks
|
142 |
+
code_blocks = re.findall(r"```(?:python)?\n(.*?)```", text, re.DOTALL)
|
143 |
+
if code_blocks:
|
144 |
+
return code_blocks[0]
|
145 |
+
return None
|
146 |
+
|
147 |
+
def remove_code_blocks(text):
|
148 |
+
# Remove code blocks from text
|
149 |
+
return re.sub(r"```(?:python)?\n.*?```", '', text, flags=re.DOTALL).strip()
|
150 |
+
|
151 |
+
def execute_code_safely(code):
|
152 |
+
"""
|
153 |
+
Executes the given code safely and returns the output.
|
154 |
+
WARNING: Executing arbitrary code can be dangerous.
|
155 |
+
This function uses restricted built-ins and a try-except block to handle errors.
|
156 |
+
In a production environment, consider using a secure sandbox or a code execution service.
|
157 |
+
"""
|
158 |
+
# Restrict built-ins
|
159 |
+
safe_globals = {
|
160 |
+
"__builtins__": {
|
161 |
+
'abs': abs,
|
162 |
+
'all': all,
|
163 |
+
'any': any,
|
164 |
+
'len': len,
|
165 |
+
'max': max,
|
166 |
+
'min': min,
|
167 |
+
'sum': sum,
|
168 |
+
'range': range,
|
169 |
+
'print': print,
|
170 |
+
'str': str,
|
171 |
+
'int': int,
|
172 |
+
'float': float,
|
173 |
+
'bool': bool,
|
174 |
+
'list': list,
|
175 |
+
'dict': dict,
|
176 |
+
'set': set,
|
177 |
+
'tuple': tuple,
|
178 |
+
'enumerate': enumerate,
|
179 |
+
'zip': zip,
|
180 |
+
'math': __import__('math'),
|
181 |
+
'datetime': __import__('datetime'),
|
182 |
+
}
|
183 |
+
}
|
184 |
+
safe_locals = {}
|
185 |
+
|
186 |
+
# Capture output
|
187 |
+
old_stdout = sys.stdout
|
188 |
+
redirected_output = sys.stdout = io.StringIO()
|
189 |
+
|
190 |
+
try:
|
191 |
+
exec(code, safe_globals, safe_locals)
|
192 |
+
except Exception as e:
|
193 |
+
output = f"Error executing code:\n{traceback.format_exc()}"
|
194 |
+
else:
|
195 |
+
output = redirected_output.getvalue()
|
196 |
+
finally:
|
197 |
+
sys.stdout = old_stdout
|
198 |
+
|
199 |
+
return output.strip()
|
200 |
+
|
201 |
+
if __name__ == "__main__":
|
202 |
+
main()
|
203 |
+
|
204 |
+
# requirements.txt
|
205 |
+
|
requirements.txt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
streamlit
|
2 |
+
groq
|
3 |
+
python-dotenv
|