computer-use-ootb-internal 0.0.95.post2__tar.gz → 0.0.96__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/PKG-INFO +1 -1
  2. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/pyproject.toml +1 -1
  3. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/app_teachmode.py +420 -387
  4. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/.gitignore +0 -0
  5. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/README.md +0 -0
  6. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/app_teachmode_gradio.py +0 -0
  7. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/animation/click_animation.py +0 -0
  8. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/animation/icons8-select-cursor-transparent-96.gif +0 -0
  9. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/executor/teachmode_executor.py +0 -0
  10. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/__init__.py +0 -0
  11. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/gui_capture.py +0 -0
  12. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/gui_parser.py +0 -0
  13. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/icon_detection/icon_detection.py +0 -0
  14. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/panel_recognition/llm_panel_recognize.py +0 -0
  15. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/test_capture.py +0 -0
  16. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/uia_parser.py +0 -0
  17. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/simple_parser/utils.py +0 -0
  18. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/uia_tools/__init__.py +0 -0
  19. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/uia_tools/screenshot_cli.py +0 -0
  20. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/gui_parser/uia_tools/screenshot_service.py +0 -0
  21. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/llm_utils/llm_utils.py +0 -0
  22. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/llm_utils/oai.py +0 -0
  23. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/llm_utils/run_litellm.py +0 -0
  24. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/llm_utils/run_llm.py +0 -0
  25. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/vlm_utils/__init__.py +0 -0
  26. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/gui_agent/vlm_utils/run_vlm.py +0 -0
  27. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/__init__.py +0 -0
  28. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/aws_request.py +0 -0
  29. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/base.py +0 -0
  30. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/bash.py +0 -0
  31. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/collection.py +0 -0
  32. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/colorful_text.py +0 -0
  33. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/computer.py +0 -0
  34. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/computer_marbot.py +0 -0
  35. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/edit.py +0 -0
  36. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/run.py +0 -0
  37. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/computer_use_demo/tools/screen_capture.py +0 -0
  38. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/dependency_check.py +0 -0
  39. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/example_websocket_js.html +0 -0
  40. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/requirements-lite.txt +0 -0
  41. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/run_teachmode_ootb_args.py +0 -0
  42. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/service_teachmode.py +0 -0
  43. {computer_use_ootb_internal-0.0.95.post2 → computer_use_ootb_internal-0.0.96}/src/computer_use_ootb_internal/service_teachmode_test.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: computer-use-ootb-internal
3
- Version: 0.0.95.post2
3
+ Version: 0.0.96
4
4
  Summary: Computer Use OOTB
5
5
  Author-email: Siyuan Hu <siyuan.hu.sg@gmail.com>
6
6
  Requires-Python: >=3.11
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "computer-use-ootb-internal"
7
- version = "0.0.95.post2"
7
+ version = "0.0.96"
8
8
  description = "Computer Use OOTB"
9
9
  authors = [{ name = "Siyuan Hu", email = "siyuan.hu.sg@gmail.com" }]
10
10
  requires-python = ">=3.11"
@@ -1,387 +1,420 @@
1
- import argparse
2
- import time
3
- import json
4
- import threading
5
- from fastapi import FastAPI, Request
6
- from fastapi.responses import JSONResponse
7
- from fastapi.middleware.cors import CORSMiddleware
8
- from screeninfo import get_monitors
9
- from computer_use_ootb_internal.computer_use_demo.tools.computer import get_screen_details
10
- from computer_use_ootb_internal.run_teachmode_ootb_args import simple_teachmode_sampling_loop
11
-
12
- app = FastAPI()
13
-
14
- # Add CORS middleware to allow requests from the frontend
15
- app.add_middleware(
16
- CORSMiddleware,
17
- allow_origins=["*"],
18
- allow_credentials=True,
19
- allow_methods=["*"],
20
- allow_headers=["*"],
21
- )
22
-
23
- # Rate limiter for API endpoints
24
- class RateLimiter:
25
- def __init__(self, interval_seconds=2):
26
- self.interval = interval_seconds
27
- self.last_request_time = {}
28
- self.lock = threading.Lock()
29
-
30
- def allow_request(self, endpoint):
31
- with self.lock:
32
- current_time = time.time()
33
- # Priority endpoints always allowed
34
- if endpoint in ["/update_params", "/update_message"]:
35
- return True
36
-
37
- # For other endpoints, apply rate limiting
38
- if endpoint not in self.last_request_time:
39
- self.last_request_time[endpoint] = current_time
40
- return True
41
-
42
- elapsed = current_time - self.last_request_time[endpoint]
43
- if elapsed < self.interval:
44
- return False
45
-
46
- self.last_request_time[endpoint] = current_time
47
- return True
48
-
49
- class SharedState:
50
- def __init__(self, args):
51
- self.args = args
52
- self.task_updated = False
53
- self.chatbot_messages = []
54
- # Store all state-related data here
55
- self.model = args.model
56
- self.task = getattr(args, 'task', "")
57
- self.selected_screen = args.selected_screen
58
- self.user_id = args.user_id
59
- self.trace_id = args.trace_id
60
- self.api_keys = args.api_keys
61
- self.server_url = args.server_url
62
- self.message_queue = []
63
- self.is_processing = False
64
- self.should_stop = False
65
- self.is_paused = False
66
- # Add a new event to better control stopping
67
- self.stop_event = threading.Event()
68
- # Add a reference to the processing thread
69
- self.processing_thread = None
70
-
71
- shared_state = None
72
- rate_limiter = RateLimiter(interval_seconds=2)
73
-
74
- @app.post("/update_params")
75
- async def update_parameters(request: Request):
76
- data = await request.json()
77
-
78
- if 'task' not in data:
79
- return JSONResponse(
80
- content={"status": "error", "message": "Missing required field: task"},
81
- status_code=400
82
- )
83
-
84
- shared_state.args = argparse.Namespace(**data)
85
- shared_state.task_updated = True
86
-
87
- # Update shared state when parameters change
88
- shared_state.model = getattr(shared_state.args, 'model', "teach-mode-gpt-4o")
89
- shared_state.task = getattr(shared_state.args, 'task', "Create a claim on the SAP system, using Receipt.pdf as attachment.")
90
- shared_state.selected_screen = getattr(shared_state.args, 'selected_screen', 0)
91
- shared_state.user_id = getattr(shared_state.args, 'user_id', "a_test")
92
- shared_state.trace_id = getattr(shared_state.args, 'trace_id', "jess_4")
93
- shared_state.api_keys = getattr(shared_state.args, 'api_keys', "sk-proj-1234567890")
94
- shared_state.server_url = getattr(shared_state.args, 'server_url', "http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com/generate_action")
95
-
96
- return JSONResponse(
97
- content={"status": "success", "message": "Parameters updated", "new_args": vars(shared_state.args)},
98
- status_code=200
99
- )
100
-
101
- @app.post("/update_message")
102
- async def update_message(request: Request):
103
- data = await request.json()
104
-
105
- if 'message' not in data:
106
- return JSONResponse(
107
- content={"status": "error", "message": "Missing required field: message"},
108
- status_code=400
109
- )
110
-
111
- message = data['message']
112
- shared_state.chatbot_messages.append({"role": "user", "content": message})
113
- shared_state.task = message
114
- shared_state.args.task = message
115
-
116
- # Reset stop event before starting
117
- shared_state.stop_event.clear()
118
-
119
- # Start processing if not already running
120
- if not shared_state.is_processing:
121
- # Create and store the thread
122
- shared_state.processing_thread = threading.Thread(target=process_input, daemon=True)
123
- shared_state.processing_thread.start()
124
-
125
- return JSONResponse(
126
- content={"status": "success", "message": "Message received", "task": shared_state.task},
127
- status_code=200
128
- )
129
-
130
- @app.get("/get_messages")
131
- async def get_messages(request: Request):
132
- # Apply rate limiting
133
- if not rate_limiter.allow_request(request.url.path):
134
- return JSONResponse(
135
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
136
- status_code=429
137
- )
138
-
139
- # Return all messages in the queue and clear it
140
- messages = shared_state.message_queue.copy()
141
- shared_state.message_queue = []
142
-
143
- return JSONResponse(
144
- content={"status": "success", "messages": messages},
145
- status_code=200
146
- )
147
-
148
- @app.get("/get_screens")
149
- async def get_screens(request: Request):
150
- # Apply rate limiting
151
- if not rate_limiter.allow_request(request.url.path):
152
- return JSONResponse(
153
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
154
- status_code=429
155
- )
156
-
157
- screen_options, primary_index = get_screen_details()
158
-
159
- return JSONResponse(
160
- content={"status": "success", "screens": screen_options, "primary_index": primary_index},
161
- status_code=200
162
- )
163
-
164
- @app.post("/stop_processing")
165
- async def stop_processing(request: Request):
166
- # Apply rate limiting
167
- if not rate_limiter.allow_request(request.url.path):
168
- return JSONResponse(
169
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
170
- status_code=429
171
- )
172
-
173
- if shared_state.is_processing:
174
- # Set both flags to ensure stopping the current task
175
- shared_state.should_stop = True
176
- shared_state.stop_event.set()
177
-
178
- # Send an immediate message to the queue to inform the user
179
- stop_initiated_msg = {"role": "assistant", "content": f"Stopping task '{shared_state.task}'..."}
180
- shared_state.message_queue.append(stop_initiated_msg)
181
-
182
- return JSONResponse(
183
- content={"status": "success", "message": "Task is being stopped, server will remain available for new tasks"},
184
- status_code=200
185
- )
186
- else:
187
- return JSONResponse(
188
- content={"status": "error", "message": "No active processing to stop"},
189
- status_code=400
190
- )
191
-
192
- @app.post("/toggle_pause")
193
- async def toggle_pause(request: Request):
194
- # Apply rate limiting
195
- if not rate_limiter.allow_request(request.url.path):
196
- return JSONResponse(
197
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
198
- status_code=429
199
- )
200
-
201
- if not shared_state.is_processing:
202
- return JSONResponse(
203
- content={"status": "error", "message": "No active processing to pause/resume"},
204
- status_code=400
205
- )
206
-
207
- # Toggle the pause state
208
- shared_state.is_paused = not shared_state.is_paused
209
- current_state = shared_state.is_paused
210
-
211
- print(f"Toggled pause state to: {current_state}")
212
-
213
- status_message = "paused" if current_state else "resumed"
214
-
215
- # Add a message to the queue to inform the user
216
- if current_state:
217
- message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been paused. Click Continue to resume."}
218
- else:
219
- message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been resumed."}
220
-
221
- shared_state.chatbot_messages.append(message)
222
- shared_state.message_queue.append(message)
223
-
224
- return JSONResponse(
225
- content={
226
- "status": "success",
227
- "message": f"Processing {status_message}",
228
- "is_paused": current_state
229
- },
230
- status_code=200
231
- )
232
-
233
- @app.get("/status")
234
- async def get_status(request: Request):
235
- # Apply rate limiting
236
- if not rate_limiter.allow_request(request.url.path):
237
- return JSONResponse(
238
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
239
- status_code=429
240
- )
241
-
242
- print(f"Status check - Processing: {shared_state.is_processing}, Paused: {shared_state.is_paused}")
243
- return JSONResponse(
244
- content={
245
- "status": "success",
246
- "is_processing": shared_state.is_processing,
247
- "is_paused": shared_state.is_paused
248
- },
249
- status_code=200
250
- )
251
-
252
- def process_input():
253
- shared_state.is_processing = True
254
- shared_state.should_stop = False
255
- shared_state.is_paused = False
256
- shared_state.stop_event.clear() # Ensure stop event is cleared at the start
257
-
258
- print(f"start sampling loop: {shared_state.chatbot_messages}")
259
- print(f"shared_state.args before sampling loop: {shared_state.args}")
260
-
261
-
262
- try:
263
- # Get the generator for the sampling loop
264
- sampling_loop = simple_teachmode_sampling_loop(
265
- model=shared_state.model,
266
- task=shared_state.task,
267
- selected_screen=shared_state.selected_screen,
268
- user_id=shared_state.user_id,
269
- trace_id=shared_state.trace_id,
270
- api_keys=shared_state.api_keys,
271
- server_url=shared_state.server_url,
272
- )
273
-
274
- # Process messages from the sampling loop
275
- for loop_msg in sampling_loop:
276
- # Check stop condition more frequently
277
- if shared_state.should_stop or shared_state.stop_event.is_set():
278
- print("Processing stopped by user")
279
- break
280
-
281
- # Check if paused and wait while paused
282
- while shared_state.is_paused and not shared_state.should_stop and not shared_state.stop_event.is_set():
283
- print(f"Processing paused at: {time.strftime('%H:%M:%S')}")
284
- # Wait a short time and check stop condition regularly
285
- for _ in range(5): # Check 5 times per second
286
- if shared_state.should_stop or shared_state.stop_event.is_set():
287
- break
288
- time.sleep(0.2)
289
-
290
- # Check again after pause loop
291
- if shared_state.should_stop or shared_state.stop_event.is_set():
292
- print("Processing stopped while paused or resuming")
293
- break
294
-
295
- # Process the message
296
- if loop_msg.startswith('<img'):
297
- message = {"role": "user", "content": loop_msg}
298
- else:
299
- message = {"role": "assistant", "content": loop_msg}
300
-
301
- shared_state.chatbot_messages.append(message)
302
- shared_state.message_queue.append(message)
303
-
304
- # Short sleep to allow stop signals to be processed
305
- for _ in range(5): # Check 5 times per second
306
- if shared_state.should_stop or shared_state.stop_event.is_set():
307
- print("Processing stopped during sleep")
308
- break
309
- time.sleep(0.1)
310
-
311
- if shared_state.should_stop or shared_state.stop_event.is_set():
312
- break
313
-
314
- except Exception as e:
315
- # Handle any exceptions in the processing loop
316
- error_msg = f"Error during task processing: {str(e)}"
317
- print(error_msg)
318
- error_message = {"role": "assistant", "content": error_msg}
319
- shared_state.message_queue.append(error_message)
320
-
321
- finally:
322
- # Handle completion or interruption
323
- if shared_state.should_stop or shared_state.stop_event.is_set():
324
- stop_msg = f"Task '{shared_state.task}' was stopped. Ready for new tasks."
325
- final_message = {"role": "assistant", "content": stop_msg}
326
- else:
327
- complete_msg = f"Task '{shared_state.task}' completed. Thanks for using Teachmode-OOTB."
328
- final_message = {"role": "assistant", "content": complete_msg}
329
-
330
- shared_state.chatbot_messages.append(final_message)
331
- shared_state.message_queue.append(final_message)
332
-
333
- # Reset all state flags to allow for new tasks
334
- shared_state.is_processing = False
335
- shared_state.should_stop = False
336
- shared_state.is_paused = False
337
- shared_state.stop_event.clear()
338
- print("Processing completed, ready for new tasks")
339
-
340
- def main():
341
- global app, shared_state, rate_limiter
342
-
343
- parser = argparse.ArgumentParser(
344
- description="Run a synchronous sampling loop for assistant/tool interactions in teach-mode."
345
- )
346
- parser.add_argument("--model", default="teach-mode-gpt-4o")
347
- parser.add_argument("--task", default="Create a claim on the SAP system, using Receipt.pdf as attachment.")
348
- parser.add_argument("--selected_screen", type=int, default=0)
349
- parser.add_argument("--user_id", default="star_rail_dev")
350
- parser.add_argument("--trace_id", default="scroll")
351
- parser.add_argument("--api_key_file", default="api_key.json")
352
- parser.add_argument("--api_keys", default="")
353
- parser.add_argument(
354
- "--server_url",
355
- default="http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com/generate_action",
356
- help="Server URL for the session"
357
- )
358
-
359
- args = parser.parse_args()
360
- shared_state = SharedState(args)
361
- rate_limiter = RateLimiter(interval_seconds=2)
362
-
363
- import uvicorn
364
- import platform
365
- import os
366
-
367
- # Default port
368
- port = 7888
369
-
370
- # Determine port based on Windows username
371
- if platform.system() == "Windows":
372
- username = os.environ["USERNAME"].lower()
373
- if username == "altair":
374
- port = 14000
375
- elif username.startswith("guest") and username[5:].isdigit():
376
- num = int(username[5:])
377
- if 1 <= num <= 10:
378
- port = 14000 + num
379
- else:
380
- port = 7888
381
- else:
382
- port = 7888
383
-
384
- uvicorn.run(app, host="0.0.0.0", port=port)
385
-
386
- if __name__ == "__main__":
387
- main()
1
+ import argparse
2
+ import time
3
+ import json
4
+ from datetime import datetime
5
+ import threading
6
+ import requests
7
+ from fastapi import FastAPI, Request
8
+ from fastapi.responses import JSONResponse
9
+ from fastapi.middleware.cors import CORSMiddleware
10
+ from screeninfo import get_monitors
11
+ from computer_use_ootb_internal.computer_use_demo.tools.computer import get_screen_details
12
+ from computer_use_ootb_internal.run_teachmode_ootb_args import simple_teachmode_sampling_loop
13
+
14
+ app = FastAPI()
15
+
16
+ # Add CORS middleware to allow requests from the frontend
17
+ app.add_middleware(
18
+ CORSMiddleware,
19
+ allow_origins=["*"],
20
+ allow_credentials=True,
21
+ allow_methods=["*"],
22
+ allow_headers=["*"],
23
+ )
24
+
25
+ # Rate limiter for API endpoints
26
+ class RateLimiter:
27
+ def __init__(self, interval_seconds=2):
28
+ self.interval = interval_seconds
29
+ self.last_request_time = {}
30
+ self.lock = threading.Lock()
31
+
32
+ def allow_request(self, endpoint):
33
+ with self.lock:
34
+ current_time = time.time()
35
+ # Priority endpoints always allowed
36
+ if endpoint in ["/update_params", "/update_message"]:
37
+ return True
38
+
39
+ # For other endpoints, apply rate limiting
40
+ if endpoint not in self.last_request_time:
41
+ self.last_request_time[endpoint] = current_time
42
+ return True
43
+
44
+ elapsed = current_time - self.last_request_time[endpoint]
45
+ if elapsed < self.interval:
46
+ return False
47
+
48
+ self.last_request_time[endpoint] = current_time
49
+ return True
50
+
51
+
52
+ def log_ootb_request(server_url, ootb_request_type, data):
53
+ logging_data = {
54
+ "type": ootb_request_type,
55
+ "data": data,
56
+ "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
57
+ }
58
+ if not server_url.endswith("/update_ootb_logging"):
59
+ server_logging_url = server_url + "/update_ootb_logging"
60
+ else:
61
+ server_logging_url = server_url
62
+ requests.post(server_logging_url, json=logging_data)
63
+
64
+
65
+ class SharedState:
66
+ def __init__(self, args):
67
+ self.args = args
68
+ self.task_updated = False
69
+ self.chatbot_messages = []
70
+ # Store all state-related data here
71
+ self.model = args.model
72
+ self.task = getattr(args, 'task', "")
73
+ self.selected_screen = args.selected_screen
74
+ self.user_id = args.user_id
75
+ self.trace_id = args.trace_id
76
+ self.api_keys = args.api_keys
77
+ self.server_url = args.server_url
78
+ self.message_queue = []
79
+ self.is_processing = False
80
+ self.should_stop = False
81
+ self.is_paused = False
82
+ # Add a new event to better control stopping
83
+ self.stop_event = threading.Event()
84
+ # Add a reference to the processing thread
85
+ self.processing_thread = None
86
+
87
+ shared_state = None
88
+ rate_limiter = RateLimiter(interval_seconds=2)
89
+
90
+ @app.post("/update_params")
91
+ async def update_parameters(request: Request):
92
+ data = await request.json()
93
+
94
+ if 'task' not in data:
95
+ return JSONResponse(
96
+ content={"status": "error", "message": "Missing required field: task"},
97
+ status_code=400
98
+ )
99
+
100
+ shared_state.args = argparse.Namespace(**data)
101
+ shared_state.task_updated = True
102
+
103
+ # Update shared state when parameters change
104
+ shared_state.model = getattr(shared_state.args, 'model', "teach-mode-gpt-4o")
105
+ shared_state.task = getattr(shared_state.args, 'task', "Create a claim on the SAP system, using Receipt.pdf as attachment.")
106
+ shared_state.selected_screen = getattr(shared_state.args, 'selected_screen', 0)
107
+ shared_state.user_id = getattr(shared_state.args, 'user_id', "a_test")
108
+ shared_state.trace_id = getattr(shared_state.args, 'trace_id', "jess_4")
109
+ shared_state.api_keys = getattr(shared_state.args, 'api_keys', "sk-proj-1234567890")
110
+ shared_state.server_url = getattr(shared_state.args, 'server_url', "http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com")
111
+
112
+ log_ootb_request(shared_state.server_url, "update_params", data)
113
+
114
+ return JSONResponse(
115
+ content={"status": "success", "message": "Parameters updated", "new_args": vars(shared_state.args)},
116
+ status_code=200
117
+ )
118
+
119
+ @app.post("/update_message")
120
+ async def update_message(request: Request):
121
+ data = await request.json()
122
+
123
+ if 'message' not in data:
124
+ return JSONResponse(
125
+ content={"status": "error", "message": "Missing required field: message"},
126
+ status_code=400
127
+ )
128
+
129
+ log_ootb_request(shared_state.server_url, "update_message", data)
130
+
131
+ message = data['message']
132
+ shared_state.chatbot_messages.append({"role": "user", "content": message})
133
+ shared_state.task = message
134
+ shared_state.args.task = message
135
+
136
+ # Reset stop event before starting
137
+ shared_state.stop_event.clear()
138
+
139
+ # Start processing if not already running
140
+ if not shared_state.is_processing:
141
+ # Create and store the thread
142
+ shared_state.processing_thread = threading.Thread(target=process_input, daemon=True)
143
+ shared_state.processing_thread.start()
144
+
145
+ return JSONResponse(
146
+ content={"status": "success", "message": "Message received", "task": shared_state.task},
147
+ status_code=200
148
+ )
149
+
150
+ @app.get("/get_messages")
151
+ async def get_messages(request: Request):
152
+ # Apply rate limiting
153
+ if not rate_limiter.allow_request(request.url.path):
154
+ return JSONResponse(
155
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
156
+ status_code=429
157
+ )
158
+
159
+ log_ootb_request(shared_state.server_url, "get_messages", {})
160
+
161
+ # Return all messages in the queue and clear it
162
+ messages = shared_state.message_queue.copy()
163
+ shared_state.message_queue = []
164
+
165
+ return JSONResponse(
166
+ content={"status": "success", "messages": messages},
167
+ status_code=200
168
+ )
169
+
170
+ @app.get("/get_screens")
171
+ async def get_screens(request: Request):
172
+ # Apply rate limiting
173
+ if not rate_limiter.allow_request(request.url.path):
174
+ return JSONResponse(
175
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
176
+ status_code=429
177
+ )
178
+
179
+ log_ootb_request(shared_state.server_url, "get_screens", {})
180
+
181
+ screen_options, primary_index = get_screen_details()
182
+
183
+ return JSONResponse(
184
+ content={"status": "success", "screens": screen_options, "primary_index": primary_index},
185
+ status_code=200
186
+ )
187
+
188
+ @app.post("/stop_processing")
189
+ async def stop_processing(request: Request):
190
+ # Apply rate limiting
191
+ if not rate_limiter.allow_request(request.url.path):
192
+ return JSONResponse(
193
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
194
+ status_code=429
195
+ )
196
+
197
+ log_ootb_request(shared_state.server_url, "stop_processing", {})
198
+
199
+ if shared_state.is_processing:
200
+ # Set both flags to ensure stopping the current task
201
+ shared_state.should_stop = True
202
+ shared_state.stop_event.set()
203
+
204
+ # Send an immediate message to the queue to inform the user
205
+ stop_initiated_msg = {"role": "assistant", "content": f"Stopping task '{shared_state.task}'..."}
206
+ shared_state.message_queue.append(stop_initiated_msg)
207
+
208
+ return JSONResponse(
209
+ content={"status": "success", "message": "Task is being stopped, server will remain available for new tasks"},
210
+ status_code=200
211
+ )
212
+ else:
213
+ return JSONResponse(
214
+ content={"status": "error", "message": "No active processing to stop"},
215
+ status_code=400
216
+ )
217
+
218
+ @app.post("/toggle_pause")
219
+ async def toggle_pause(request: Request):
220
+ # Apply rate limiting
221
+ if not rate_limiter.allow_request(request.url.path):
222
+ return JSONResponse(
223
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
224
+ status_code=429
225
+ )
226
+
227
+ log_ootb_request(shared_state.server_url, "toggle_pause", {})
228
+
229
+ if not shared_state.is_processing:
230
+ return JSONResponse(
231
+ content={"status": "error", "message": "No active processing to pause/resume"},
232
+ status_code=400
233
+ )
234
+
235
+ # Toggle the pause state
236
+ shared_state.is_paused = not shared_state.is_paused
237
+ current_state = shared_state.is_paused
238
+
239
+ print(f"Toggled pause state to: {current_state}")
240
+
241
+ status_message = "paused" if current_state else "resumed"
242
+
243
+ # Add a message to the queue to inform the user
244
+ if current_state:
245
+ message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been paused. Click Continue to resume."}
246
+ else:
247
+ message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been resumed."}
248
+
249
+ shared_state.chatbot_messages.append(message)
250
+ shared_state.message_queue.append(message)
251
+
252
+ return JSONResponse(
253
+ content={
254
+ "status": "success",
255
+ "message": f"Processing {status_message}",
256
+ "is_paused": current_state
257
+ },
258
+ status_code=200
259
+ )
260
+
261
+ @app.get("/status")
262
+ async def get_status(request: Request):
263
+ # Apply rate limiting
264
+ if not rate_limiter.allow_request(request.url.path):
265
+ return JSONResponse(
266
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
267
+ status_code=429
268
+ )
269
+
270
+ log_ootb_request(shared_state.server_url, "get_status", {})
271
+
272
+ print(f"Status check - Processing: {shared_state.is_processing}, Paused: {shared_state.is_paused}")
273
+ return JSONResponse(
274
+ content={
275
+ "status": "success",
276
+ "is_processing": shared_state.is_processing,
277
+ "is_paused": shared_state.is_paused
278
+ },
279
+ status_code=200
280
+ )
281
+
282
+ def process_input():
283
+ shared_state.is_processing = True
284
+ shared_state.should_stop = False
285
+ shared_state.is_paused = False
286
+ shared_state.stop_event.clear() # Ensure stop event is cleared at the start
287
+
288
+ print(f"start sampling loop: {shared_state.chatbot_messages}")
289
+ print(f"shared_state.args before sampling loop: {shared_state.args}")
290
+
291
+
292
+ try:
293
+ # Get the generator for the sampling loop
294
+ sampling_loop = simple_teachmode_sampling_loop(
295
+ model=shared_state.model,
296
+ task=shared_state.task,
297
+ selected_screen=shared_state.selected_screen,
298
+ user_id=shared_state.user_id,
299
+ trace_id=shared_state.trace_id,
300
+ api_keys=shared_state.api_keys,
301
+ server_url=shared_state.server_url,
302
+ )
303
+
304
+ # Process messages from the sampling loop
305
+ for loop_msg in sampling_loop:
306
+ # Check stop condition more frequently
307
+ if shared_state.should_stop or shared_state.stop_event.is_set():
308
+ print("Processing stopped by user")
309
+ break
310
+
311
+ # Check if paused and wait while paused
312
+ while shared_state.is_paused and not shared_state.should_stop and not shared_state.stop_event.is_set():
313
+ print(f"Processing paused at: {time.strftime('%H:%M:%S')}")
314
+ # Wait a short time and check stop condition regularly
315
+ for _ in range(5): # Check 5 times per second
316
+ if shared_state.should_stop or shared_state.stop_event.is_set():
317
+ break
318
+ time.sleep(0.2)
319
+
320
+ # Check again after pause loop
321
+ if shared_state.should_stop or shared_state.stop_event.is_set():
322
+ print("Processing stopped while paused or resuming")
323
+ break
324
+
325
+ # Process the message
326
+ if loop_msg.startswith('<img'):
327
+ message = {"role": "user", "content": loop_msg}
328
+ else:
329
+ message = {"role": "assistant", "content": loop_msg}
330
+
331
+ shared_state.chatbot_messages.append(message)
332
+ shared_state.message_queue.append(message)
333
+
334
+ # Short sleep to allow stop signals to be processed
335
+ for _ in range(5): # Check 5 times per second
336
+ if shared_state.should_stop or shared_state.stop_event.is_set():
337
+ print("Processing stopped during sleep")
338
+ break
339
+ time.sleep(0.1)
340
+
341
+ if shared_state.should_stop or shared_state.stop_event.is_set():
342
+ break
343
+
344
+ except Exception as e:
345
+ # Handle any exceptions in the processing loop
346
+ error_msg = f"Error during task processing: {str(e)}"
347
+ print(error_msg)
348
+ error_message = {"role": "assistant", "content": error_msg}
349
+ shared_state.message_queue.append(error_message)
350
+
351
+ finally:
352
+ # Handle completion or interruption
353
+ if shared_state.should_stop or shared_state.stop_event.is_set():
354
+ stop_msg = f"Task '{shared_state.task}' was stopped. Ready for new tasks."
355
+ final_message = {"role": "assistant", "content": stop_msg}
356
+ else:
357
+ complete_msg = f"Task '{shared_state.task}' completed. Thanks for using Teachmode-OOTB."
358
+ final_message = {"role": "assistant", "content": complete_msg}
359
+
360
+ shared_state.chatbot_messages.append(final_message)
361
+ shared_state.message_queue.append(final_message)
362
+
363
+ # Reset all state flags to allow for new tasks
364
+ shared_state.is_processing = False
365
+ shared_state.should_stop = False
366
+ shared_state.is_paused = False
367
+ shared_state.stop_event.clear()
368
+ print("Processing completed, ready for new tasks")
369
+
370
+ def main():
371
+ global app, shared_state, rate_limiter
372
+
373
+ parser = argparse.ArgumentParser(
374
+ description="Run a synchronous sampling loop for assistant/tool interactions in teach-mode."
375
+ )
376
+ parser.add_argument("--model", default="teach-mode-gpt-4o")
377
+ parser.add_argument("--task", default="Create a claim on the SAP system, using Receipt.pdf as attachment.")
378
+ parser.add_argument("--selected_screen", type=int, default=0)
379
+ parser.add_argument("--user_id", default="star_rail_dev")
380
+ parser.add_argument("--trace_id", default="scroll")
381
+ parser.add_argument("--api_key_file", default="api_key.json")
382
+ parser.add_argument("--api_keys", default="")
383
+ parser.add_argument(
384
+ "--server_url",
385
+ default="http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com",
386
+ help="Server URL for the session"
387
+ )
388
+
389
+ args = parser.parse_args()
390
+ shared_state = SharedState(args)
391
+ rate_limiter = RateLimiter(interval_seconds=2)
392
+
393
+ import uvicorn
394
+ import platform
395
+ import os
396
+
397
+ # Default port
398
+ port = 7888
399
+
400
+ # Determine port based on Windows username
401
+ if platform.system() == "Windows":
402
+ username = os.environ["USERNAME"].lower()
403
+ if username == "altair":
404
+ port = 14000
405
+ elif username.startswith("guest") and username[5:].isdigit():
406
+ num = int(username[5:])
407
+ if 1 <= num <= 10:
408
+ port = 14000 + num
409
+ else:
410
+ port = 7888
411
+ else:
412
+ port = 7888
413
+
414
+ uvicorn.run(app, host="0.0.0.0", port=port)
415
+
416
+ if __name__ == "__main__":
417
+ # main()
418
+
419
+ # Test log_ootb_request
420
+ log_ootb_request("http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com", "test_request", {"message": "Test message"})