computer-use-ootb-internal 0.0.105__py3-none-any.whl → 0.0.106__py3-none-any.whl

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.
@@ -1,491 +1,491 @@
1
- import argparse
2
- import time
3
- import json
4
- from datetime import datetime
5
- import threading
6
- import requests
7
- import platform # Add platform import
8
- import subprocess # Add subprocess import
9
- import pyautogui # Add pyautogui import
10
- import webbrowser # Add webbrowser import
11
- import os # Import os for path joining
12
- from fastapi import FastAPI, Request
13
- from fastapi.responses import JSONResponse
14
- from fastapi.middleware.cors import CORSMiddleware
15
- from screeninfo import get_monitors
16
- from computer_use_ootb_internal.computer_use_demo.tools.computer import get_screen_details
17
- from computer_use_ootb_internal.run_teachmode_ootb_args import simple_teachmode_sampling_loop
18
-
19
- app = FastAPI()
20
-
21
- # Add CORS middleware to allow requests from the frontend
22
- app.add_middleware(
23
- CORSMiddleware,
24
- allow_origins=["*"],
25
- allow_credentials=True,
26
- allow_methods=["*"],
27
- allow_headers=["*"],
28
- )
29
-
30
- # Rate limiter for API endpoints
31
- class RateLimiter:
32
- def __init__(self, interval_seconds=2):
33
- self.interval = interval_seconds
34
- self.last_request_time = {}
35
- self.lock = threading.Lock()
36
-
37
- def allow_request(self, endpoint):
38
- with self.lock:
39
- current_time = time.time()
40
- # Priority endpoints always allowed
41
- if endpoint in ["/update_params", "/update_message"]:
42
- return True
43
-
44
- # For other endpoints, apply rate limiting
45
- if endpoint not in self.last_request_time:
46
- self.last_request_time[endpoint] = current_time
47
- return True
48
-
49
- elapsed = current_time - self.last_request_time[endpoint]
50
- if elapsed < self.interval:
51
- return False
52
-
53
- self.last_request_time[endpoint] = current_time
54
- return True
55
-
56
-
57
- def log_ootb_request(server_url, ootb_request_type, data):
58
- logging_data = {
59
- "type": ootb_request_type,
60
- "data": data,
61
- "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
62
- }
63
- if not server_url.endswith("/update_ootb_logging"):
64
- server_logging_url = server_url + "/update_ootb_logging"
65
- else:
66
- server_logging_url = server_url
67
- requests.post(server_logging_url, json=logging_data)
68
-
69
-
70
- class SharedState:
71
- def __init__(self, args):
72
- self.args = args
73
- self.task_updated = False
74
- self.chatbot_messages = []
75
- # Store all state-related data here
76
- self.model = args.model
77
- self.task = getattr(args, 'task', "")
78
- self.selected_screen = args.selected_screen
79
- self.user_id = args.user_id
80
- self.trace_id = args.trace_id
81
- self.api_keys = args.api_keys
82
- self.server_url = args.server_url
83
- self.message_queue = []
84
- self.is_processing = False
85
- self.should_stop = False
86
- self.is_paused = False
87
- # Add a new event to better control stopping
88
- self.stop_event = threading.Event()
89
- # Add a reference to the processing thread
90
- self.processing_thread = None
91
-
92
- shared_state = None
93
- rate_limiter = RateLimiter(interval_seconds=2)
94
-
95
- # Add the new prepare_environment function here
96
- def prepare_environment(state):
97
- """Prepares the environment before starting the main processing loop, e.g., opening specific apps."""
98
- if platform.system() == "Windows":
99
- # Assuming Star Rail mode is indicated by user_id containing "star_rail"
100
- # You might need to adjust this condition based on the actual logic in run_teachmode_args
101
- is_star_rail = "star_rail" in state.user_id.lower() or \
102
- "star_rail" in state.trace_id.lower() or \
103
- "hero_case" in state.trace_id.lower()
104
-
105
- if is_star_rail:
106
- print("Star Rail mode detected on Windows. Opening Edge browser...")
107
- url = "https://sr.mihoyo.com/cloud/#/"
108
- browser_opened = False
109
- try:
110
- # Use only webbrowser.open
111
- print(f"Attempting to open {url} using webbrowser.open()...")
112
- if webbrowser.open(url):
113
- print(f"Successfully requested browser to open {url} via webbrowser.open().")
114
- browser_opened = True
115
- else:
116
- print("webbrowser.open() returned False, indicating potential failure.")
117
-
118
- if not browser_opened:
119
- print("ERROR: Failed to confirm browser opening via webbrowser.open().")
120
- # Still proceed to click attempt
121
-
122
- # Add pyautogui click after attempting to open the browser
123
- print("Proceeding with pyautogui actions...")
124
- time.sleep(5) # Wait time for the browser to load
125
-
126
- # Print detected screen size
127
- screen_width, screen_height = pyautogui.size()
128
- print(f"Detected screen size: {screen_width}x{screen_height}")
129
-
130
- click_x = int(screen_width * (1036 / 1280))
131
- click_y = int(screen_height * (500 / 720))
132
- print(f"Calculated click coordinates: ({click_x}, {click_y})")
133
-
134
- # Disable failsafe before clicking
135
- pyautogui.FAILSAFE = False
136
- print("PyAutoGUI failsafe temporarily disabled.")
137
-
138
- print(f"Clicking at coordinates: ({click_x}, {click_y})")
139
- pyautogui.click(click_x, click_y)
140
- time.sleep(2)
141
- pyautogui.click(click_x, click_y)
142
-
143
- # Re-enable failsafe (optional, as script might end anyway)
144
- # pyautogui.FAILSAFE = True
145
- # print("PyAutoGUI failsafe re-enabled.")
146
-
147
- except Exception as e:
148
- print(f"Error during environment preparation (browser/click): {e}")
149
- finally:
150
- # Ensure failsafe is re-enabled if an error occurs after disabling it
151
- pyautogui.FAILSAFE = True
152
- print("PyAutoGUI failsafe re-enabled.")
153
- else:
154
- # Placeholder for potential preparations on other OS or non-Star Rail modes
155
- print("Environment preparation: No specific actions required for this OS/mode.")
156
-
157
-
158
- @app.post("/update_params")
159
- async def update_parameters(request: Request):
160
- data = await request.json()
161
-
162
- if 'task' not in data:
163
- return JSONResponse(
164
- content={"status": "error", "message": "Missing required field: task"},
165
- status_code=400
166
- )
167
-
168
- shared_state.args = argparse.Namespace(**data)
169
- shared_state.task_updated = True
170
-
171
- # Update shared state when parameters change
172
- shared_state.model = getattr(shared_state.args, 'model', "teach-mode-gpt-4o")
173
- shared_state.task = getattr(shared_state.args, 'task', "Create a claim on the SAP system, using Receipt.pdf as attachment.")
174
- shared_state.selected_screen = getattr(shared_state.args, 'selected_screen', 0)
175
- shared_state.user_id = getattr(shared_state.args, 'user_id', "hero_cases")
176
- shared_state.trace_id = getattr(shared_state.args, 'trace_id', "build_scroll_combat")
177
- shared_state.api_keys = getattr(shared_state.args, 'api_keys', "sk-proj-1234567890")
178
- shared_state.server_url = getattr(shared_state.args, 'server_url', "http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com")
179
-
180
- log_ootb_request(shared_state.server_url, "update_params", data)
181
-
182
- # Call the preparation function here, after parameters are updated
183
- prepare_environment(shared_state)
184
-
185
- return JSONResponse(
186
- content={"status": "success", "message": "Parameters updated", "new_args": vars(shared_state.args)},
187
- status_code=200
188
- )
189
-
190
- @app.post("/update_message")
191
- async def update_message(request: Request):
192
- data = await request.json()
193
-
194
- if 'message' not in data:
195
- return JSONResponse(
196
- content={"status": "error", "message": "Missing required field: message"},
197
- status_code=400
198
- )
199
-
200
- log_ootb_request(shared_state.server_url, "update_message", data)
201
-
202
- message = data['message']
203
- shared_state.chatbot_messages.append({"role": "user", "content": message})
204
- shared_state.task = message
205
- shared_state.args.task = message
206
-
207
- # Reset stop event before starting
208
- shared_state.stop_event.clear()
209
-
210
- # Start processing if not already running
211
- if not shared_state.is_processing:
212
- # Create and store the thread
213
- shared_state.processing_thread = threading.Thread(target=process_input, daemon=True)
214
- shared_state.processing_thread.start()
215
-
216
- return JSONResponse(
217
- content={"status": "success", "message": "Message received", "task": shared_state.task},
218
- status_code=200
219
- )
220
-
221
- @app.get("/get_messages")
222
- async def get_messages(request: Request):
223
- # Apply rate limiting
224
- if not rate_limiter.allow_request(request.url.path):
225
- return JSONResponse(
226
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
227
- status_code=429
228
- )
229
-
230
- log_ootb_request(shared_state.server_url, "get_messages", {})
231
-
232
- # Return all messages in the queue and clear it
233
- messages = shared_state.message_queue.copy()
234
- shared_state.message_queue = []
235
-
236
- return JSONResponse(
237
- content={"status": "success", "messages": messages},
238
- status_code=200
239
- )
240
-
241
- @app.get("/get_screens")
242
- async def get_screens(request: Request):
243
- # Apply rate limiting
244
- if not rate_limiter.allow_request(request.url.path):
245
- return JSONResponse(
246
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
247
- status_code=429
248
- )
249
-
250
- log_ootb_request(shared_state.server_url, "get_screens", {})
251
-
252
- screen_options, primary_index = get_screen_details()
253
-
254
- return JSONResponse(
255
- content={"status": "success", "screens": screen_options, "primary_index": primary_index},
256
- status_code=200
257
- )
258
-
259
- @app.post("/stop_processing")
260
- async def stop_processing(request: Request):
261
- # Apply rate limiting
262
- if not rate_limiter.allow_request(request.url.path):
263
- return JSONResponse(
264
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
265
- status_code=429
266
- )
267
-
268
- log_ootb_request(shared_state.server_url, "stop_processing", {})
269
-
270
- if shared_state.is_processing:
271
- # Set both flags to ensure stopping the current task
272
- shared_state.should_stop = True
273
- shared_state.stop_event.set()
274
-
275
- # Send an immediate message to the queue to inform the user
276
- stop_initiated_msg = {"role": "assistant", "content": f"Stopping task '{shared_state.task}'..."}
277
- shared_state.message_queue.append(stop_initiated_msg)
278
-
279
- return JSONResponse(
280
- content={"status": "success", "message": "Task is being stopped, server will remain available for new tasks"},
281
- status_code=200
282
- )
283
- else:
284
- return JSONResponse(
285
- content={"status": "error", "message": "No active processing to stop"},
286
- status_code=400
287
- )
288
-
289
- @app.post("/toggle_pause")
290
- async def toggle_pause(request: Request):
291
- # Apply rate limiting
292
- if not rate_limiter.allow_request(request.url.path):
293
- return JSONResponse(
294
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
295
- status_code=429
296
- )
297
-
298
- log_ootb_request(shared_state.server_url, "toggle_pause", {})
299
-
300
- if not shared_state.is_processing:
301
- return JSONResponse(
302
- content={"status": "error", "message": "No active processing to pause/resume"},
303
- status_code=400
304
- )
305
-
306
- # Toggle the pause state
307
- shared_state.is_paused = not shared_state.is_paused
308
- current_state = shared_state.is_paused
309
-
310
- print(f"Toggled pause state to: {current_state}")
311
-
312
- status_message = "paused" if current_state else "resumed"
313
-
314
- # Add a message to the queue to inform the user
315
- if current_state:
316
- message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been paused. Click Continue to resume."}
317
- else:
318
- message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been resumed."}
319
-
320
- shared_state.chatbot_messages.append(message)
321
- shared_state.message_queue.append(message)
322
-
323
- return JSONResponse(
324
- content={
325
- "status": "success",
326
- "message": f"Processing {status_message}",
327
- "is_paused": current_state
328
- },
329
- status_code=200
330
- )
331
-
332
- @app.get("/status")
333
- async def get_status(request: Request):
334
- # Apply rate limiting
335
- if not rate_limiter.allow_request(request.url.path):
336
- return JSONResponse(
337
- content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
338
- status_code=429
339
- )
340
-
341
- log_ootb_request(shared_state.server_url, "get_status", {})
342
-
343
- print(f"Status check - Processing: {shared_state.is_processing}, Paused: {shared_state.is_paused}")
344
- return JSONResponse(
345
- content={
346
- "status": "success",
347
- "is_processing": shared_state.is_processing,
348
- "is_paused": shared_state.is_paused
349
- },
350
- status_code=200
351
- )
352
-
353
- def process_input():
354
- shared_state.is_processing = True
355
- shared_state.should_stop = False
356
- shared_state.is_paused = False
357
- shared_state.stop_event.clear() # Ensure stop event is cleared at the start
358
-
359
- print(f"start sampling loop: {shared_state.chatbot_messages}")
360
- print(f"shared_state.args before sampling loop: {shared_state.args}")
361
-
362
-
363
- try:
364
- # Get the generator for the sampling loop
365
- sampling_loop = simple_teachmode_sampling_loop(
366
- model=shared_state.model,
367
- task=shared_state.task,
368
- selected_screen=shared_state.selected_screen,
369
- user_id=shared_state.user_id,
370
- trace_id=shared_state.trace_id,
371
- api_keys=shared_state.api_keys,
372
- server_url=shared_state.server_url,
373
- )
374
-
375
- # Process messages from the sampling loop
376
- for loop_msg in sampling_loop:
377
- # Check stop condition more frequently
378
- if shared_state.should_stop or shared_state.stop_event.is_set():
379
- print("Processing stopped by user")
380
- break
381
-
382
- # Check if paused and wait while paused
383
- while shared_state.is_paused and not shared_state.should_stop and not shared_state.stop_event.is_set():
384
- print(f"Processing paused at: {time.strftime('%H:%M:%S')}")
385
- # Wait a short time and check stop condition regularly
386
- for _ in range(5): # Check 5 times per second
387
- if shared_state.should_stop or shared_state.stop_event.is_set():
388
- break
389
- time.sleep(0.2)
390
-
391
- # Check again after pause loop
392
- if shared_state.should_stop or shared_state.stop_event.is_set():
393
- print("Processing stopped while paused or resuming")
394
- break
395
-
396
- # Process the message
397
- if loop_msg.startswith('<img'):
398
- message = {"role": "user", "content": loop_msg}
399
- else:
400
- message = {"role": "assistant", "content": loop_msg}
401
-
402
- shared_state.chatbot_messages.append(message)
403
- shared_state.message_queue.append(message)
404
-
405
- # Short sleep to allow stop signals to be processed
406
- for _ in range(5): # Check 5 times per second
407
- if shared_state.should_stop or shared_state.stop_event.is_set():
408
- print("Processing stopped during sleep")
409
- break
410
- time.sleep(0.1)
411
-
412
- if shared_state.should_stop or shared_state.stop_event.is_set():
413
- break
414
-
415
- except Exception as e:
416
- # Handle any exceptions in the processing loop
417
- error_msg = f"Error during task processing: {str(e)}"
418
- print(error_msg)
419
- error_message = {"role": "assistant", "content": error_msg}
420
- shared_state.message_queue.append(error_message)
421
-
422
- finally:
423
- # Handle completion or interruption
424
- if shared_state.should_stop or shared_state.stop_event.is_set():
425
- stop_msg = f"Task '{shared_state.task}' was stopped. Ready for new tasks."
426
- final_message = {"role": "assistant", "content": stop_msg}
427
- else:
428
- complete_msg = f"Task '{shared_state.task}' completed. Thanks for using Teachmode-OOTB."
429
- final_message = {"role": "assistant", "content": complete_msg}
430
-
431
- shared_state.chatbot_messages.append(final_message)
432
- shared_state.message_queue.append(final_message)
433
-
434
- # Reset all state flags to allow for new tasks
435
- shared_state.is_processing = False
436
- shared_state.should_stop = False
437
- shared_state.is_paused = False
438
- shared_state.stop_event.clear()
439
- print("Processing completed, ready for new tasks")
440
-
441
- def main():
442
- global app, shared_state, rate_limiter
443
-
444
- parser = argparse.ArgumentParser(
445
- description="Run a synchronous sampling loop for assistant/tool interactions in teach-mode."
446
- )
447
- parser.add_argument("--model", default="teach-mode-gpt-4o")
448
- parser.add_argument("--task", default="Create a claim on the SAP system, using Receipt.pdf as attachment.")
449
- parser.add_argument("--selected_screen", type=int, default=0)
450
- parser.add_argument("--user_id", default="star_rail_dev")
451
- parser.add_argument("--trace_id", default="scroll")
452
- parser.add_argument("--api_key_file", default="api_key.json")
453
- parser.add_argument("--api_keys", default="")
454
- parser.add_argument(
455
- "--server_url",
456
- default="http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com",
457
- help="Server URL for the session"
458
- )
459
-
460
- args = parser.parse_args()
461
- shared_state = SharedState(args)
462
- rate_limiter = RateLimiter(interval_seconds=2)
463
-
464
- import uvicorn
465
- import platform
466
- import os
467
-
468
- # Default port
469
- port = 7888
470
-
471
- # Determine port based on Windows username
472
- if platform.system() == "Windows":
473
- username = os.environ["USERNAME"].lower()
474
- if username == "altair":
475
- port = 14000
476
- elif username.startswith("guest") and username[5:].isdigit():
477
- num = int(username[5:])
478
- if 1 <= num <= 10:
479
- port = 14000 + num
480
- else:
481
- port = 7888
482
- else:
483
- port = 7888
484
-
485
- uvicorn.run(app, host="0.0.0.0", port=port)
486
-
487
- if __name__ == "__main__":
488
- # main()
489
-
490
- # Test log_ootb_request
1
+ import argparse
2
+ import time
3
+ import json
4
+ from datetime import datetime
5
+ import threading
6
+ import requests
7
+ import platform # Add platform import
8
+ import subprocess # Add subprocess import
9
+ import pyautogui # Add pyautogui import
10
+ import webbrowser # Add webbrowser import
11
+ import os # Import os for path joining
12
+ from fastapi import FastAPI, Request
13
+ from fastapi.responses import JSONResponse
14
+ from fastapi.middleware.cors import CORSMiddleware
15
+ from screeninfo import get_monitors
16
+ from computer_use_ootb_internal.computer_use_demo.tools.computer import get_screen_details
17
+ from computer_use_ootb_internal.run_teachmode_ootb_args import simple_teachmode_sampling_loop
18
+
19
+ app = FastAPI()
20
+
21
+ # Add CORS middleware to allow requests from the frontend
22
+ app.add_middleware(
23
+ CORSMiddleware,
24
+ allow_origins=["*"],
25
+ allow_credentials=True,
26
+ allow_methods=["*"],
27
+ allow_headers=["*"],
28
+ )
29
+
30
+ # Rate limiter for API endpoints
31
+ class RateLimiter:
32
+ def __init__(self, interval_seconds=2):
33
+ self.interval = interval_seconds
34
+ self.last_request_time = {}
35
+ self.lock = threading.Lock()
36
+
37
+ def allow_request(self, endpoint):
38
+ with self.lock:
39
+ current_time = time.time()
40
+ # Priority endpoints always allowed
41
+ if endpoint in ["/update_params", "/update_message"]:
42
+ return True
43
+
44
+ # For other endpoints, apply rate limiting
45
+ if endpoint not in self.last_request_time:
46
+ self.last_request_time[endpoint] = current_time
47
+ return True
48
+
49
+ elapsed = current_time - self.last_request_time[endpoint]
50
+ if elapsed < self.interval:
51
+ return False
52
+
53
+ self.last_request_time[endpoint] = current_time
54
+ return True
55
+
56
+
57
+ def log_ootb_request(server_url, ootb_request_type, data):
58
+ logging_data = {
59
+ "type": ootb_request_type,
60
+ "data": data,
61
+ "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
62
+ }
63
+ if not server_url.endswith("/update_ootb_logging"):
64
+ server_logging_url = server_url + "/update_ootb_logging"
65
+ else:
66
+ server_logging_url = server_url
67
+ requests.post(server_logging_url, json=logging_data)
68
+
69
+
70
+ class SharedState:
71
+ def __init__(self, args):
72
+ self.args = args
73
+ self.task_updated = False
74
+ self.chatbot_messages = []
75
+ # Store all state-related data here
76
+ self.model = args.model
77
+ self.task = getattr(args, 'task', "")
78
+ self.selected_screen = args.selected_screen
79
+ self.user_id = args.user_id
80
+ self.trace_id = args.trace_id
81
+ self.api_keys = args.api_keys
82
+ self.server_url = args.server_url
83
+ self.message_queue = []
84
+ self.is_processing = False
85
+ self.should_stop = False
86
+ self.is_paused = False
87
+ # Add a new event to better control stopping
88
+ self.stop_event = threading.Event()
89
+ # Add a reference to the processing thread
90
+ self.processing_thread = None
91
+
92
+ shared_state = None
93
+ rate_limiter = RateLimiter(interval_seconds=2)
94
+
95
+ # Add the new prepare_environment function here
96
+ def prepare_environment(state):
97
+ """Prepares the environment before starting the main processing loop, e.g., opening specific apps."""
98
+ if platform.system() == "Windows":
99
+ # Assuming Star Rail mode is indicated by user_id containing "star_rail"
100
+ # You might need to adjust this condition based on the actual logic in run_teachmode_args
101
+ is_star_rail = "star_rail" in state.user_id.lower() or \
102
+ "star_rail" in state.trace_id.lower() or \
103
+ "hero_case" in state.trace_id.lower()
104
+
105
+ if is_star_rail:
106
+ print("Star Rail mode detected on Windows. Opening Edge browser...")
107
+ url = "https://sr.mihoyo.com/cloud/#/"
108
+ browser_opened = False
109
+ try:
110
+ # Use only webbrowser.open
111
+ print(f"Attempting to open {url} using webbrowser.open()...")
112
+ if webbrowser.open(url):
113
+ print(f"Successfully requested browser to open {url} via webbrowser.open().")
114
+ browser_opened = True
115
+ else:
116
+ print("webbrowser.open() returned False, indicating potential failure.")
117
+
118
+ if not browser_opened:
119
+ print("ERROR: Failed to confirm browser opening via webbrowser.open().")
120
+ # Still proceed to click attempt
121
+
122
+ # Add pyautogui click after attempting to open the browser
123
+ print("Proceeding with pyautogui actions...")
124
+ time.sleep(5) # Wait time for the browser to load
125
+
126
+ # Print detected screen size
127
+ screen_width, screen_height = pyautogui.size()
128
+ print(f"Detected screen size: {screen_width}x{screen_height}")
129
+
130
+ click_x = int(screen_width * (1036 / 1280))
131
+ click_y = int(screen_height * (500 / 720))
132
+ print(f"Calculated click coordinates: ({click_x}, {click_y})")
133
+
134
+ # Disable failsafe before clicking
135
+ pyautogui.FAILSAFE = False
136
+ print("PyAutoGUI failsafe temporarily disabled.")
137
+
138
+ print(f"Clicking at coordinates: ({click_x}, {click_y})")
139
+ pyautogui.click(click_x, click_y)
140
+ time.sleep(2)
141
+ pyautogui.click(click_x, click_y)
142
+
143
+ # Re-enable failsafe (optional, as script might end anyway)
144
+ # pyautogui.FAILSAFE = True
145
+ # print("PyAutoGUI failsafe re-enabled.")
146
+
147
+ except Exception as e:
148
+ print(f"Error during environment preparation (browser/click): {e}")
149
+ finally:
150
+ # Ensure failsafe is re-enabled if an error occurs after disabling it
151
+ pyautogui.FAILSAFE = True
152
+ print("PyAutoGUI failsafe re-enabled.")
153
+ else:
154
+ # Placeholder for potential preparations on other OS or non-Star Rail modes
155
+ print("Environment preparation: No specific actions required for this OS/mode.")
156
+
157
+
158
+ @app.post("/update_params")
159
+ async def update_parameters(request: Request):
160
+ data = await request.json()
161
+
162
+ if 'task' not in data:
163
+ return JSONResponse(
164
+ content={"status": "error", "message": "Missing required field: task"},
165
+ status_code=400
166
+ )
167
+
168
+ shared_state.args = argparse.Namespace(**data)
169
+ shared_state.task_updated = True
170
+
171
+ # Update shared state when parameters change
172
+ shared_state.model = getattr(shared_state.args, 'model', "teach-mode-gpt-4o")
173
+ shared_state.task = getattr(shared_state.args, 'task', "Create a claim on the SAP system, using Receipt.pdf as attachment.")
174
+ shared_state.selected_screen = getattr(shared_state.args, 'selected_screen', 0)
175
+ shared_state.user_id = getattr(shared_state.args, 'user_id', "hero_cases")
176
+ shared_state.trace_id = getattr(shared_state.args, 'trace_id', "build_scroll_combat")
177
+ shared_state.api_keys = getattr(shared_state.args, 'api_keys', "sk-proj-1234567890")
178
+ shared_state.server_url = getattr(shared_state.args, 'server_url', "http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com")
179
+
180
+ log_ootb_request(shared_state.server_url, "update_params", data)
181
+
182
+ # Call the preparation function here, after parameters are updated
183
+ prepare_environment(shared_state)
184
+
185
+ return JSONResponse(
186
+ content={"status": "success", "message": "Parameters updated", "new_args": vars(shared_state.args)},
187
+ status_code=200
188
+ )
189
+
190
+ @app.post("/update_message")
191
+ async def update_message(request: Request):
192
+ data = await request.json()
193
+
194
+ if 'message' not in data:
195
+ return JSONResponse(
196
+ content={"status": "error", "message": "Missing required field: message"},
197
+ status_code=400
198
+ )
199
+
200
+ log_ootb_request(shared_state.server_url, "update_message", data)
201
+
202
+ message = data['message']
203
+ shared_state.chatbot_messages.append({"role": "user", "content": message})
204
+ shared_state.task = message
205
+ shared_state.args.task = message
206
+
207
+ # Reset stop event before starting
208
+ shared_state.stop_event.clear()
209
+
210
+ # Start processing if not already running
211
+ if not shared_state.is_processing:
212
+ # Create and store the thread
213
+ shared_state.processing_thread = threading.Thread(target=process_input, daemon=True)
214
+ shared_state.processing_thread.start()
215
+
216
+ return JSONResponse(
217
+ content={"status": "success", "message": "Message received", "task": shared_state.task},
218
+ status_code=200
219
+ )
220
+
221
+ @app.get("/get_messages")
222
+ async def get_messages(request: Request):
223
+ # Apply rate limiting
224
+ if not rate_limiter.allow_request(request.url.path):
225
+ return JSONResponse(
226
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
227
+ status_code=429
228
+ )
229
+
230
+ log_ootb_request(shared_state.server_url, "get_messages", {})
231
+
232
+ # Return all messages in the queue and clear it
233
+ messages = shared_state.message_queue.copy()
234
+ shared_state.message_queue = []
235
+
236
+ return JSONResponse(
237
+ content={"status": "success", "messages": messages},
238
+ status_code=200
239
+ )
240
+
241
+ @app.get("/get_screens")
242
+ async def get_screens(request: Request):
243
+ # Apply rate limiting
244
+ if not rate_limiter.allow_request(request.url.path):
245
+ return JSONResponse(
246
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
247
+ status_code=429
248
+ )
249
+
250
+ log_ootb_request(shared_state.server_url, "get_screens", {})
251
+
252
+ screen_options, primary_index = get_screen_details()
253
+
254
+ return JSONResponse(
255
+ content={"status": "success", "screens": screen_options, "primary_index": primary_index},
256
+ status_code=200
257
+ )
258
+
259
+ @app.post("/stop_processing")
260
+ async def stop_processing(request: Request):
261
+ # Apply rate limiting
262
+ if not rate_limiter.allow_request(request.url.path):
263
+ return JSONResponse(
264
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
265
+ status_code=429
266
+ )
267
+
268
+ log_ootb_request(shared_state.server_url, "stop_processing", {})
269
+
270
+ if shared_state.is_processing:
271
+ # Set both flags to ensure stopping the current task
272
+ shared_state.should_stop = True
273
+ shared_state.stop_event.set()
274
+
275
+ # Send an immediate message to the queue to inform the user
276
+ stop_initiated_msg = {"role": "assistant", "content": f"Stopping task '{shared_state.task}'..."}
277
+ shared_state.message_queue.append(stop_initiated_msg)
278
+
279
+ return JSONResponse(
280
+ content={"status": "success", "message": "Task is being stopped, server will remain available for new tasks"},
281
+ status_code=200
282
+ )
283
+ else:
284
+ return JSONResponse(
285
+ content={"status": "error", "message": "No active processing to stop"},
286
+ status_code=400
287
+ )
288
+
289
+ @app.post("/toggle_pause")
290
+ async def toggle_pause(request: Request):
291
+ # Apply rate limiting
292
+ if not rate_limiter.allow_request(request.url.path):
293
+ return JSONResponse(
294
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
295
+ status_code=429
296
+ )
297
+
298
+ log_ootb_request(shared_state.server_url, "toggle_pause", {})
299
+
300
+ if not shared_state.is_processing:
301
+ return JSONResponse(
302
+ content={"status": "error", "message": "No active processing to pause/resume"},
303
+ status_code=400
304
+ )
305
+
306
+ # Toggle the pause state
307
+ shared_state.is_paused = not shared_state.is_paused
308
+ current_state = shared_state.is_paused
309
+
310
+ print(f"Toggled pause state to: {current_state}")
311
+
312
+ status_message = "paused" if current_state else "resumed"
313
+
314
+ # Add a message to the queue to inform the user
315
+ if current_state:
316
+ message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been paused. Click Continue to resume."}
317
+ else:
318
+ message = {"role": "assistant", "content": f"Task '{shared_state.task}' has been resumed."}
319
+
320
+ shared_state.chatbot_messages.append(message)
321
+ shared_state.message_queue.append(message)
322
+
323
+ return JSONResponse(
324
+ content={
325
+ "status": "success",
326
+ "message": f"Processing {status_message}",
327
+ "is_paused": current_state
328
+ },
329
+ status_code=200
330
+ )
331
+
332
+ @app.get("/status")
333
+ async def get_status(request: Request):
334
+ # Apply rate limiting
335
+ if not rate_limiter.allow_request(request.url.path):
336
+ return JSONResponse(
337
+ content={"status": "error", "message": "Rate limit exceeded. Try again after 2 seconds."},
338
+ status_code=429
339
+ )
340
+
341
+ log_ootb_request(shared_state.server_url, "get_status", {})
342
+
343
+ print(f"Status check - Processing: {shared_state.is_processing}, Paused: {shared_state.is_paused}")
344
+ return JSONResponse(
345
+ content={
346
+ "status": "success",
347
+ "is_processing": shared_state.is_processing,
348
+ "is_paused": shared_state.is_paused
349
+ },
350
+ status_code=200
351
+ )
352
+
353
+ def process_input():
354
+ shared_state.is_processing = True
355
+ shared_state.should_stop = False
356
+ shared_state.is_paused = False
357
+ shared_state.stop_event.clear() # Ensure stop event is cleared at the start
358
+
359
+ print(f"start sampling loop: {shared_state.chatbot_messages}")
360
+ print(f"shared_state.args before sampling loop: {shared_state.args}")
361
+
362
+
363
+ try:
364
+ # Get the generator for the sampling loop
365
+ sampling_loop = simple_teachmode_sampling_loop(
366
+ model=shared_state.model,
367
+ task=shared_state.task,
368
+ selected_screen=shared_state.selected_screen,
369
+ user_id=shared_state.user_id,
370
+ trace_id=shared_state.trace_id,
371
+ api_keys=shared_state.api_keys,
372
+ server_url=shared_state.server_url,
373
+ )
374
+
375
+ # Process messages from the sampling loop
376
+ for loop_msg in sampling_loop:
377
+ # Check stop condition more frequently
378
+ if shared_state.should_stop or shared_state.stop_event.is_set():
379
+ print("Processing stopped by user")
380
+ break
381
+
382
+ # Check if paused and wait while paused
383
+ while shared_state.is_paused and not shared_state.should_stop and not shared_state.stop_event.is_set():
384
+ print(f"Processing paused at: {time.strftime('%H:%M:%S')}")
385
+ # Wait a short time and check stop condition regularly
386
+ for _ in range(5): # Check 5 times per second
387
+ if shared_state.should_stop or shared_state.stop_event.is_set():
388
+ break
389
+ time.sleep(0.2)
390
+
391
+ # Check again after pause loop
392
+ if shared_state.should_stop or shared_state.stop_event.is_set():
393
+ print("Processing stopped while paused or resuming")
394
+ break
395
+
396
+ # Process the message
397
+ if loop_msg.startswith('<img'):
398
+ message = {"role": "user", "content": loop_msg}
399
+ else:
400
+ message = {"role": "assistant", "content": loop_msg}
401
+
402
+ shared_state.chatbot_messages.append(message)
403
+ shared_state.message_queue.append(message)
404
+
405
+ # Short sleep to allow stop signals to be processed
406
+ for _ in range(5): # Check 5 times per second
407
+ if shared_state.should_stop or shared_state.stop_event.is_set():
408
+ print("Processing stopped during sleep")
409
+ break
410
+ time.sleep(0.1)
411
+
412
+ if shared_state.should_stop or shared_state.stop_event.is_set():
413
+ break
414
+
415
+ except Exception as e:
416
+ # Handle any exceptions in the processing loop
417
+ error_msg = f"Error during task processing: {str(e)}"
418
+ print(error_msg)
419
+ error_message = {"role": "assistant", "content": error_msg}
420
+ shared_state.message_queue.append(error_message)
421
+
422
+ finally:
423
+ # Handle completion or interruption
424
+ if shared_state.should_stop or shared_state.stop_event.is_set():
425
+ stop_msg = f"Task '{shared_state.task}' was stopped. Ready for new tasks."
426
+ final_message = {"role": "assistant", "content": stop_msg}
427
+ else:
428
+ complete_msg = f"Task '{shared_state.task}' completed. Thanks for using Teachmode-OOTB."
429
+ final_message = {"role": "assistant", "content": complete_msg}
430
+
431
+ shared_state.chatbot_messages.append(final_message)
432
+ shared_state.message_queue.append(final_message)
433
+
434
+ # Reset all state flags to allow for new tasks
435
+ shared_state.is_processing = False
436
+ shared_state.should_stop = False
437
+ shared_state.is_paused = False
438
+ shared_state.stop_event.clear()
439
+ print("Processing completed, ready for new tasks")
440
+
441
+ def main():
442
+ global app, shared_state, rate_limiter
443
+
444
+ parser = argparse.ArgumentParser(
445
+ description="Run a synchronous sampling loop for assistant/tool interactions in teach-mode."
446
+ )
447
+ parser.add_argument("--model", default="teach-mode-gpt-4o")
448
+ parser.add_argument("--task", default="Create a claim on the SAP system, using Receipt.pdf as attachment.")
449
+ parser.add_argument("--selected_screen", type=int, default=0)
450
+ parser.add_argument("--user_id", default="star_rail_dev")
451
+ parser.add_argument("--trace_id", default="scroll")
452
+ parser.add_argument("--api_key_file", default="api_key.json")
453
+ parser.add_argument("--api_keys", default="")
454
+ parser.add_argument(
455
+ "--server_url",
456
+ default="http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com",
457
+ help="Server URL for the session"
458
+ )
459
+
460
+ args = parser.parse_args()
461
+ shared_state = SharedState(args)
462
+ rate_limiter = RateLimiter(interval_seconds=2)
463
+
464
+ import uvicorn
465
+ import platform
466
+ import os
467
+
468
+ # Default port
469
+ port = 7888
470
+
471
+ # Determine port based on Windows username
472
+ if platform.system() == "Windows":
473
+ username = os.environ["USERNAME"].lower()
474
+ if username == "altair":
475
+ port = 14000
476
+ elif username.startswith("guest") and username[5:].isdigit():
477
+ num = int(username[5:])
478
+ if 1 <= num <= 10:
479
+ port = 14000 + num
480
+ else:
481
+ port = 7888
482
+ else:
483
+ port = 7888
484
+
485
+ uvicorn.run(app, host="0.0.0.0", port=port)
486
+
487
+ if __name__ == "__main__":
488
+ # main()
489
+
490
+ # Test log_ootb_request
491
491
  log_ootb_request("http://ec2-44-234-43-86.us-west-2.compute.amazonaws.com", "test_request", {"message": "Test message"})