Divax commited on
Commit
5f1c192
·
1 Parent(s): d677138

files added

Browse files
Files changed (2) hide show
  1. app.py +205 -0
  2. 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