refactor: optimize inference with async locking, prompt engineering, and input post-processing while adding comprehensive testing suites

app.py

@@ -55,11 +55,15 @@ if not MOCK_MODE:
     nemotron = Llama(model_path=NEMOTRON_MODEL_PATH, n_ctx=1024)
     minicpm = Llama(model_path=MINICPM_MODEL_PATH, n_ctx=1024)
 
+llama_lock = asyncio.Lock()
 
 async def run_blocking(func, *args, **kwargs):
-    return await asyncio.to_thread(func, *args, **kwargs)
+    async with llama_lock:
+        return await asyncio.to_thread(func, *args, **kwargs)
 
-async def generate_atomic_task(goal: str, previous_failures: int, completed_count: int = 0) -> str:
+async def generate_atomic_task(goal: str, previous_failures: int, history: list = None) -> str:
+    if history is None:
+        history = []
     if MOCK_MODE:
         await asyncio.sleep(1) # Simulate inference latency
         demo_steps = [
@@ -70,20 +74,51 @@ async def generate_atomic_task(goal: str, previous_failures: int, completed_coun
         ]
         if previous_failures == 1:
             return "Move your mouse to the browser icon."
-        return demo_steps[min(completed_count, len(demo_steps) - 1)]
+        return demo_steps[min(len(history), len(demo_steps) - 1)]
             
     else:
         # REAL INFERENCE
-        prompt = f"Goal: {goal}. Failures: {previous_failures}. Output a single physical action verb under 8 words."
+        system_msg = "You are a cognitive pacemaker. Break down goals into extremely tiny, atomic physical actions under 8 words."
+        
+        prompt = f"<extra_id_0>System\n{system_msg}\n\n"
+        prompt += f"<extra_id_1>User\nGoal: write an essay\nCompleted Tasks: None\nFailures: 0\nOutput the NEXT step.\n"
+        prompt += f"<extra_id_1>Assistant\nOpen your laptop.\n"
+        prompt += f"<extra_id_1>User\nGoal: write an essay\nCompleted Tasks: - Open your laptop.\nFailures: 0\nOutput the NEXT step.\n"
+        prompt += f"<extra_id_1>Assistant\nOpen a blank document.\n"
+        prompt += f"<extra_id_1>User\nGoal: wash dishes\nCompleted Tasks: - Put dishes in sink.\nFailures: 1\nOutput the NEXT step.\n"
+        prompt += f"<extra_id_1>Assistant\nPick up a sponge.\n"
+        prompt += f"<extra_id_1>User\nGoal: {goal}\nCompleted Tasks: {history_str}\nFailures: {previous_failures}\nOutput the NEXT step.\n"
+        prompt += f"<extra_id_1>Assistant\n"
+        
         response = await run_blocking(
             nemotron,
-            prompt, 
+            prompt,
             max_tokens=20, 
-            grammar=grammar,
             temperature=0.1,
+            stop=["\n", "<extra_id_1>"]
         )
-        task_json = json.loads(response['choices'][0]['text'])
-        return task_json.get("task", "Focus on the screen.")
+        content = response['choices'][0]['text'].strip()
+        
+        # POST-PROCESSING: 4B models often ignore constraints and try to output lists.
+        # Force it to be atomic by taking the first item before a comma or " and ".
+        if "," in content:
+            content = content.split(",")[0].strip()
+        if " and " in content.lower():
+            content = content.split(" and ")[0].strip()
+            
+        content = content.strip('"\' ,.-1234567890)')
+        
+        # POST-PROCESSING: 4B models often ignore constraints and try to output lists.
+        # Force it to be atomic by taking the first item before a comma or " and ".
+        if "," in content:
+            content = content.split(",")[0].strip()
+        if " and " in content.lower():
+            content = content.split(" and ")[0].strip()
+            
+        content = content.strip('"\' ,.-1234567890)')
+        
+        print(f"NEMOTRON RAW OUT: {content}", flush=True)
+        return content or "Focus on the screen."
 
 async def apply_activation_style(task: str, style="direct") -> str:
     if MOCK_MODE:
@@ -96,9 +131,22 @@ async def apply_activation_style(task: str, style="direct") -> str:
         
     else:
         # REAL INFERENCE
-        prompt = f"Rewrite this task in a {style} tone. Keep it under 2 sentences: {task}"
-        response = await run_blocking(minicpm, prompt, max_tokens=40, temperature=0.7)
-        return response['choices'][0]['text'].strip()
+        if style == "direct":
+            return f"{task.capitalize()}."
+            
+        messages = [
+            {"role": "system", "content": f"You are a helpful assistant. Rewrite the following task into a {style} tone. Only output the final task text and nothing else."},
+            {"role": "user", "content": f"Task: {task}"}
+        ]
+        response = await run_blocking(
+            minicpm.create_chat_completion,
+            messages=messages,
+            max_tokens=40,
+            temperature=0.2
+        )
+        content = response['choices'][0]['message']['content'].strip()
+        print(f"MINICPM RAW OUT: {content}", flush=True)
+        return content
 
 @app.get("/")
 async def get():
@@ -182,7 +230,7 @@ async def websocket_endpoint(websocket: WebSocket):
             raw_task = await generate_atomic_task(
                 state["goal"],
                 state["too_hard_count"],
-                completed_count=len(state["history"]),
+                history=state["history"],
             )
             styled_task = await apply_activation_style(raw_task, state["style"])
             state["current_task"] = styled_task

print_template.py

@@ -0,0 +1,4 @@
+from llama_cpp import Llama
+import json
+nem = Llama(model_path="./models/nemotron-mini-4b.gguf", n_ctx=1024, verbose=False)
+print("TEMPLATE:", nem.metadata.get("tokenizer.chat_template"))

semantic_test.py

@@ -0,0 +1,38 @@
+import asyncio
+import websockets
+import json
+
+GOALS = [
+    "clean my messy bedroom",
+    "write a python script to parse a csv",
+    "i feel overwhelmed by my taxes",
+    "cook dinner for 4 people",
+    "learn to play the guitar",
+    "i don't know what to do",
+    "read a 500 page book",
+    "plan a vacation to japan"
+]
+
+async def semantic_test():
+    uri = "ws://127.0.0.1:7860/ws"
+    
+    print("Running Semantic Stress Test...\n")
+    
+    for goal in GOALS:
+        async with websockets.connect(uri) as ws:
+            # Action: Start
+            await ws.send(json.dumps({"action": "start", "goal": goal, "style": "direct"}))
+            resp1 = await ws.recv()
+            task1 = json.loads(resp1)['task']
+            
+            # Action: Done (Get the NEXT logical step)
+            await ws.send(json.dumps({"action": "done", "last_task": task1}))
+            resp2 = await ws.recv()
+            task2 = json.loads(resp2)['task']
+            
+            print(f"GOAL: '{goal}'")
+            print(f"  -> STEP 1: {task1}")
+            print(f"  -> STEP 2: {task2}\n")
+
+if __name__ == "__main__":
+    asyncio.run(semantic_test())

stress_test.py

@@ -0,0 +1,51 @@
+import asyncio
+import websockets
+import json
+import time
+
+async def stress_test():
+    uri = "ws://127.0.0.1:7860/ws"
+    
+    print("Connecting to WebSocket...")
+    async with websockets.connect(uri) as ws:
+        # Scenario 1: Direct style, "clean the house"
+        print("\n--- SCENARIO 1: Direct Style ---")
+        await ws.send(json.dumps({"action": "start", "goal": "clean the house", "style": "direct"}))
+        resp = await ws.recv()
+        print(f"[START] {json.loads(resp)['task']}")
+        
+        # Simulate "too_hard"
+        await ws.send(json.dumps({"action": "too_hard"}))
+        resp = await ws.recv()
+        print(f"[TOO HARD 1] {json.loads(resp)['task']}")
+        
+        # Simulate "done"
+        await ws.send(json.dumps({"action": "done", "last_task": json.loads(resp)['task']}))
+        resp = await ws.recv()
+        print(f"[DONE] {json.loads(resp)['task']}")
+        
+    async with websockets.connect(uri) as ws:
+        # Scenario 2: Calm style, "write my thesis"
+        print("\n--- SCENARIO 2: Calm Style ---")
+        await ws.send(json.dumps({"action": "start", "goal": "write my thesis", "style": "calm"}))
+        resp = await ws.recv()
+        print(f"[START] {json.loads(resp)['task']}")
+
+    async with websockets.connect(uri) as ws:
+        # Scenario 3: Encouraging style, "exercise"
+        print("\n--- SCENARIO 3: Encouraging Style ---")
+        await ws.send(json.dumps({"action": "start", "goal": "exercise", "style": "encouraging"}))
+        resp = await ws.recv()
+        print(f"[START] {json.loads(resp)['task']}")
+        
+        # Trigger circuit breaker (too hard x3)
+        for i in range(3):
+            await ws.send(json.dumps({"action": "too_hard"}))
+            resp = await ws.recv()
+            data = json.loads(resp)
+            print(f"[TOO HARD {i+1}] {data['task']} (Breaker: {data['is_breaker']})")
+
+    print("\nStress test completed successfully.")
+
+if __name__ == "__main__":
+    asyncio.run(stress_test())

test_extract.py

@@ -0,0 +1,16 @@
+import asyncio
+from llama_cpp import Llama
+
+minicpm = Llama(model_path="./models/minicpm-3-4b.gguf", n_ctx=1024, verbose=False)
+
+def extract(text):
+    messages = [
+        {"role": "system", "content": "You extract the very first physical action from the user's text. Return ONLY a single physical action under 8 words. Do not include any other words."},
+        {"role": "user", "content": f"Text: {text}"}
+    ]
+    response = minicpm.create_chat_completion(messages=messages, max_tokens=15, temperature=0.1)
+    return response['choices'][0]['message']['content'].strip()
+
+print("1:", extract("Clean bedroom floor, tidy bed, organize drawers, dust surfaces, vacuum carpet."))
+print("2:", extract("Sure, here's a python script that uses the pandas library to parse a csv file:"))
+print("3:", extract("Fail to identify goal."))

test_minicpm.py

@@ -0,0 +1,8 @@
+from llama_cpp import Llama
+minicpm = Llama(model_path="./models/minicpm-3-4b.gguf", n_ctx=1024, verbose=False)
+messages = [
+    {"role": "system", "content": "Rewrite the user's task in a direct tone. Keep it under 2 sentences. Output only the rewritten task."},
+    {"role": "user", "content": "Task: do some chores"}
+]
+response = minicpm.create_chat_completion(messages=messages, max_tokens=40, temperature=0.7)
+print("RESPONSE:", repr(response['choices'][0]['message']['content']))

test_nem_out.txt

@@ -0,0 +1,4 @@
+Error in sitecustomize; set PYTHONVERBOSE for traceback:
+ModuleNotFoundError: No module named 'wrapt'
+llama_context: n_ctx_seq (1024) < n_ctx_train (4096) -- the full capacity of the model will not be utilized
+NEMOTRON: '{\n  "task": "vacuum, dust, mop, clean bathroom, clean kitchen,'

test_nemotron.py

@@ -0,0 +1,9 @@
+import json
+from llama_cpp import Llama
+nemotron = Llama(model_path="./models/nemotron-mini-4b.gguf", n_ctx=1024, verbose=False)
+messages = [
+    {"role": "system", "content": "You are a cognitive pacemaker. Break down goals into extremely tiny, atomic physical actions under 8 words. Return ONLY a JSON object with a 'task' key."},
+    {"role": "user", "content": "Goal: clean the house. Failures: 0."}
+]
+response = nemotron.create_chat_completion(messages=messages, max_tokens=20, temperature=0.1, response_format={"type": "json_object"})
+print("NEMOTRON:", repr(response['choices'][0]['message']['content']))

test_ws.py

@@ -0,0 +1,11 @@
+import asyncio
+import websockets
+import json
+
+async def test():
+    async with websockets.connect("ws://127.0.0.1:7860/ws") as websocket:
+        await websocket.send(json.dumps({"action": "start", "goal": "clean the house", "style": "direct"}))
+        response = await websocket.recv()
+        print(f"WS RESPONSE: {response}")
+
+asyncio.run(test())