Spaces:
Running
Running
Daniel Huynh
commited on
Commit
β’
8191c52
1
Parent(s):
9578783
fix: remote exec
Browse files
app.py
CHANGED
@@ -1,22 +1,11 @@
|
|
1 |
import gradio as gr
|
2 |
-
from lavague.ActionEngine import ActionEngine
|
3 |
-
from lavague.defaults import DefaultLocalLLM, DefaultLLM
|
4 |
-
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
|
5 |
import base64
|
6 |
import requests
|
7 |
import uuid
|
8 |
-
import re
|
9 |
-
|
10 |
-
MAX_CHARS = 1500
|
11 |
|
12 |
USER_ID = str(uuid.uuid4())
|
13 |
SERVER_URL = "https://lavague.mithrilsecurity.io"
|
14 |
|
15 |
-
# Use this action_engine instead to have a local inference
|
16 |
-
# action_engine = ActionEngine(llm=DefaultLocalLLM())
|
17 |
-
|
18 |
-
action_engine = ActionEngine()
|
19 |
-
|
20 |
title = """
|
21 |
<div align="center">
|
22 |
<h1>π Welcome to LaVague</h1>
|
@@ -26,12 +15,12 @@ title = """
|
|
26 |
|
27 |
# action_engine = ActionEngine(llm, embedder)
|
28 |
|
29 |
-
def exec_code_req(url,
|
30 |
headers = {
|
31 |
"X-User-ID": USER_ID # Include the X-User-ID header for authentication
|
32 |
}
|
33 |
try:
|
34 |
-
response = requests.post(SERVER_URL + "/execute_req", json={"url": url, "
|
35 |
if response.status_code == 200:
|
36 |
return response.json()
|
37 |
else:
|
@@ -73,42 +62,21 @@ def process_url(url):
|
|
73 |
f.write(scr)
|
74 |
return "screenshot.png"
|
75 |
|
76 |
-
def
|
77 |
-
r = get_html(url_input)
|
78 |
-
state = r["html"]
|
79 |
-
query_engine = action_engine.get_query_engine(state)
|
80 |
-
response = query_engine.query(query)
|
81 |
-
source_nodes = response.get_formatted_sources(MAX_CHARS)
|
82 |
-
return response.response, source_nodes
|
83 |
-
|
84 |
-
def extract_first_python_code(markdown_text):
|
85 |
-
# Pattern to match the first ```python ``` code block
|
86 |
-
pattern = r"```python(.*?)```"
|
87 |
-
|
88 |
-
# Using re.DOTALL to make '.' match also newlines
|
89 |
-
match = re.search(pattern, markdown_text, re.DOTALL)
|
90 |
-
if match:
|
91 |
-
# Return the first matched group, which is the code inside the ```python ```
|
92 |
-
return match.group(1).strip()
|
93 |
-
else:
|
94 |
-
# Return None if no match is found
|
95 |
-
return None
|
96 |
-
|
97 |
-
|
98 |
-
def exec_code(code, source_nodes, full_code, url):
|
99 |
-
code = extract_first_python_code(code)
|
100 |
html = ""
|
101 |
try:
|
102 |
-
r = exec_code_req(url,
|
103 |
url = r["url"]
|
104 |
html = r["html"]
|
|
|
|
|
105 |
output = "Successful code execution"
|
106 |
status = """<p style="color: green; font-size: 20px; font-weight: bold;">Success!</p>"""
|
107 |
full_code += code
|
108 |
except Exception as e:
|
109 |
output = f"Error in code execution: {str(e)}"
|
110 |
status = """<p style="color: red; font-size: 20px; font-weight: bold;">Failure! Open the Debug tab for more information</p>"""
|
111 |
-
return output, code, html, status, full_code, url
|
112 |
|
113 |
def update_image_display(img, url):
|
114 |
r = send_request(url)
|
@@ -156,10 +124,8 @@ def create_demo(base_url, instructions):
|
|
156 |
# Linking components
|
157 |
url_input.submit(process_url, inputs=url_input, outputs=image_display, queue=False)
|
158 |
text_area.submit(show_processing_message, outputs=[status_html], queue=False).then(
|
159 |
-
|
160 |
-
|
161 |
-
exec_code, inputs=[code_display, source_display, full_code, url_input],
|
162 |
-
outputs=[log_display, code_display, full_html, status_html, full_code, url_input], queue=False
|
163 |
).then(
|
164 |
update_image_display, inputs=[image_display, url_input], outputs=[image_display, url_input], queue=False
|
165 |
)
|
|
|
1 |
import gradio as gr
|
|
|
|
|
|
|
2 |
import base64
|
3 |
import requests
|
4 |
import uuid
|
|
|
|
|
|
|
5 |
|
6 |
USER_ID = str(uuid.uuid4())
|
7 |
SERVER_URL = "https://lavague.mithrilsecurity.io"
|
8 |
|
|
|
|
|
|
|
|
|
|
|
9 |
title = """
|
10 |
<div align="center">
|
11 |
<h1>π Welcome to LaVague</h1>
|
|
|
15 |
|
16 |
# action_engine = ActionEngine(llm, embedder)
|
17 |
|
18 |
+
def exec_code_req(url, query):
|
19 |
headers = {
|
20 |
"X-User-ID": USER_ID # Include the X-User-ID header for authentication
|
21 |
}
|
22 |
try:
|
23 |
+
response = requests.post(SERVER_URL + "/execute_req", json={"url": url, "query": query}, headers=headers)
|
24 |
if response.status_code == 200:
|
25 |
return response.json()
|
26 |
else:
|
|
|
62 |
f.write(scr)
|
63 |
return "screenshot.png"
|
64 |
|
65 |
+
def exec_code(code, source_nodes, full_code, url, query):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
66 |
html = ""
|
67 |
try:
|
68 |
+
r = exec_code_req(url, query)
|
69 |
url = r["url"]
|
70 |
html = r["html"]
|
71 |
+
code = r["code"]
|
72 |
+
source_nodes = r["source_nodes"]
|
73 |
output = "Successful code execution"
|
74 |
status = """<p style="color: green; font-size: 20px; font-weight: bold;">Success!</p>"""
|
75 |
full_code += code
|
76 |
except Exception as e:
|
77 |
output = f"Error in code execution: {str(e)}"
|
78 |
status = """<p style="color: red; font-size: 20px; font-weight: bold;">Failure! Open the Debug tab for more information</p>"""
|
79 |
+
return output, code, html, status, full_code, url, source_nodes
|
80 |
|
81 |
def update_image_display(img, url):
|
82 |
r = send_request(url)
|
|
|
124 |
# Linking components
|
125 |
url_input.submit(process_url, inputs=url_input, outputs=image_display, queue=False)
|
126 |
text_area.submit(show_processing_message, outputs=[status_html], queue=False).then(
|
127 |
+
exec_code, inputs=[code_display, source_display, full_code, url_input, text_area],
|
128 |
+
outputs=[log_display, code_display, full_html, status_html, full_code, url_input, source_display], queue=False
|
|
|
|
|
129 |
).then(
|
130 |
update_image_display, inputs=[image_display, url_input], outputs=[image_display, url_input], queue=False
|
131 |
)
|