fix(async): fix concurrency w/ threading & asyncio

.dockerignore

@@ -16,6 +16,10 @@ requirements-dev.txt
 .idea/
 .vscode/
 *.swp
+.venv/
+venv/
+env/
+.agents/
 
 # Python cache
 __pycache__/

agent/classifier.py

@@ -12,6 +12,7 @@
 
 import logging
 import os
+import threading
 from typing import Any, Optional
 
 import torch
@@ -61,12 +62,15 @@ class SemanticClassifier:
     """
 
     _instance: Optional["SemanticClassifier"] = None
+    _init_lock: threading.Lock = threading.Lock()  # guards concurrent thread-pool callers
 
     @classmethod
     def get_instance(cls) -> "SemanticClassifier":
         if cls._instance is None:
-            logger.info("Initializing SemanticClassifier (loading model & PyTorch head)...")
-            cls._instance = cls()
+            with cls._init_lock:  # only one thread enters at a time
+                if cls._instance is None:  # re-check inside lock (double-checked locking)
+                    logger.info("Initializing SemanticClassifier (loading model & PyTorch head)...")
+                    cls._instance = cls()
         return cls._instance
 
     def __init__(self) -> None:

agent/router.py

@@ -4,10 +4,11 @@
     Layer 1a: Semantic Cache (0 tokens)
     Layer 1b: AST Math Evaluator (0 tokens, deterministic)
     Layer 2:  Weighted Scoring Classifier (~0ms, no model)
-    Layer 3:  Local SLM — Qwen2.5-3B via llama.cpp
+    Layer 3:  Local SLM — Qwen2.5 via llama.cpp
     Layer 4:  Remote Fireworks API (category-aware model + prompt compression)
 """
 
+import asyncio
 import logging
 
 from agent.ast_eval import evaluate_math_expression
@@ -78,9 +79,10 @@ async def route(self, task_id: str, prompt: str) -> str:
             return math_result
 
         # -------------------------------------------------------------------
-        # Layer 2: Weighted Classifier
+        # Layer 2: Weighted Classifier (offloaded to thread pool)
         # -------------------------------------------------------------------
-        route = classify(prompt)
+        loop = asyncio.get_running_loop()
+        route = await loop.run_in_executor(None, classify, prompt)
         logger.info("[%s] Route → %s", task_id, route)
 
         # -------------------------------------------------------------------

engines/remote_llm.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 import re
@@ -125,7 +126,13 @@ def __init__(self) -> None:
         reraise=True,
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, min=2, max=10),
-        retry=retry_if_exception_type((aiohttp.ClientResponseError, aiohttp.ClientConnectorError)),
+        retry=retry_if_exception_type(
+            (
+                aiohttp.ClientResponseError,
+                aiohttp.ClientConnectorError,
+                asyncio.TimeoutError,
+            )
+        ),
     )
     async def generate(
         self,
@@ -176,7 +183,8 @@ async def generate(
                 "temperature": temperature,
             }
 
-        async with aiohttp.ClientSession() as session:
+        timeout = aiohttp.ClientTimeout(total=20, connect=5)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
             async with session.post(endpoint, headers=headers, json=payload) as response:
                 if response.status != 200:
                     text = await response.text()

handlers/code_gen.py

@@ -1,67 +1,23 @@
 # handlers/code_gen.py
-"""Code generation handler — local-first with lean API fallback.
-
-Strategy (mirrors the AiTesting "Surgical Strike" approach):
-  easy/medium  → local Qwen 3B (0 API tokens) → validate via ast.parse
-                 → lean API fallback only if validation fails
-  hard         → direct lean API call (skip wasted local attempt)
-"""
+"""Code generation handler — direct remote API routing for speed and correctness."""
 
 import logging
 
-from engines.local_slm import LocalSLMEngine
 from engines.remote_llm import RemoteLLMEngine
 from handlers._base import load_prompt_template
-from handlers.code_utils import (
-    HARD,
-    classify_code_difficulty,
-    extract_code,
-    validate_code,
-)
+from handlers.code_utils import extract_code
 
 logger = logging.getLogger(__name__)
 
 
 class CodeGenHandler:
-    """Generates Python code with minimal API token spend."""
+    """Generates Python code using the remote model."""
 
     def __init__(self) -> None:
-        self.local = LocalSLMEngine.get_instance()
         self.remote = RemoteLLMEngine()
-        self.local_prompt = load_prompt_template("local_code_gen.txt")
         self.remote_prompt = load_prompt_template("remote_code.txt")
 
     async def handle(self, prompt: str) -> str:
-        difficulty = classify_code_difficulty(prompt)
-        logger.info("CodeGen difficulty=%s", difficulty)
-
-        # Hard tasks skip the local attempt — it almost always fails and
-        # burns wall-clock time for nothing.
-        if difficulty == HARD:
-            raw = await self.remote.generate(
-                prompt=prompt,
-                category="API_CODE",
-                system_prompt=self.remote_prompt,
-                max_tokens=500,
-            )
-            return extract_code(raw)
-
-        # Local-first for easy/medium (0 tokens)
-        local_raw = self.local.generate(
-            prompt=f"{prompt}\n\nImplement as Python code.",
-            system_prompt=self.local_prompt,
-            max_tokens=1024,
-        )
-
-        if local_raw != "__ESCALATE__":
-            code = extract_code(local_raw)
-            result = validate_code(code)
-            if result["is_valid"]:
-                logger.info("CodeGen solved locally (0 tokens)")
-                return code
-            logger.info("CodeGen local invalid: %s → API fallback", result["reason"])
-
-        # Lean one-shot API fallback
         raw = await self.remote.generate(
             prompt=prompt,
             category="API_CODE",

handlers/debug.py

@@ -1,67 +1,23 @@
 # handlers/debug.py
-"""Code debugging handler — local-first with lean API fallback.
-
-Strategy (mirrors the AiTesting "Surgical Strike" approach):
-  easy/medium  → local Qwen 3B (0 API tokens) → validate via ast.parse
-                 → lean API fallback only if validation fails
-  hard         → direct lean API call (skip wasted local attempt)
-"""
+"""Code debugging handler — direct remote API routing for speed and correctness."""
 
 import logging
 
-from engines.local_slm import LocalSLMEngine
 from engines.remote_llm import RemoteLLMEngine
 from handlers._base import load_prompt_template
-from handlers.code_utils import (
-    HARD,
-    classify_code_difficulty,
-    extract_code,
-    validate_code,
-)
+from handlers.code_utils import extract_code
 
 logger = logging.getLogger(__name__)
 
 
 class DebugHandler:
-    """Fixes buggy Python code with minimal API token spend."""
+    """Fixes buggy Python code using the remote model."""
 
     def __init__(self) -> None:
-        self.local = LocalSLMEngine.get_instance()
         self.remote = RemoteLLMEngine()
-        self.local_prompt = load_prompt_template("local_code_debug.txt")
         self.remote_prompt = load_prompt_template("remote_code.txt")
 
     async def handle(self, prompt: str) -> str:
-        difficulty = classify_code_difficulty(prompt)
-        logger.info("Debug difficulty=%s", difficulty)
-
-        # Hard tasks skip the local attempt — it almost always fails and
-        # burns wall-clock time for nothing.
-        if difficulty == HARD:
-            raw = await self.remote.generate(
-                prompt=prompt,
-                category="API_CODE",
-                system_prompt=self.remote_prompt,
-                max_tokens=400,
-            )
-            return extract_code(raw)
-
-        # Local-first for easy/medium (0 tokens)
-        local_raw = self.local.generate(
-            prompt=f"{prompt}\n\nOutput the complete corrected code.",
-            system_prompt=self.local_prompt,
-            max_tokens=1024,
-        )
-
-        if local_raw != "__ESCALATE__":
-            code = extract_code(local_raw)
-            result = validate_code(code)
-            if result["is_valid"]:
-                logger.info("Debug solved locally (0 tokens)")
-                return code
-            logger.info("Debug local invalid: %s → API fallback", result["reason"])
-
-        # Lean one-shot API fallback
         raw = await self.remote.generate(
             prompt=prompt,
             category="API_CODE",

main.py

@@ -12,6 +12,7 @@
 from dotenv import load_dotenv
 
 from agent.cache import SemanticCache
+from agent.classifier import SemanticClassifier
 from agent.router import AgentRouter
 from agent.schemas import Task
 from agent.watchdog import Watchdog
@@ -27,25 +28,43 @@
 INPUT_PATH = os.environ.get("INPUT_PATH", "/input/tasks.json")
 OUTPUT_PATH = os.environ.get("OUTPUT_PATH", "/output/results.json")
 
-# Global state
-_cache = SemanticCache()
-_router = AgentRouter(cache=_cache)
+# Global state — lightweight; expensive models are deferred into main()
+_cache: SemanticCache | None = None
+_router: AgentRouter | None = None
 _completed: list[dict[str, str]] = []
-_lock = asyncio.Lock()
+_lock: asyncio.Lock | None = None
 
 
 def _get_results() -> list[dict[str, str]]:
     return list(_completed)
 
 
 async def _process(task: Task) -> None:
+    assert _router is not None and _lock is not None
     answer = await _router.route(task.task_id, task.prompt)
     async with _lock:
         _completed.append({"task_id": task.task_id, "answer": answer})
     logger.info("Done: %s", task.task_id)
 
 
 async def main() -> None:
+    global _cache, _router, _lock
+
+    # -----------------------------------------------------------------------
+    # Initialization block — runs INSIDE asyncio.run(), after loop is live.
+    # Guards allow tests to pre-inject mocks before calling main().
+    # -----------------------------------------------------------------------
+    if _router is None:
+        logger.info("Initializing agent components (GGUF + classifier)...")
+        _cache = SemanticCache()
+        _router = AgentRouter(cache=_cache)
+        # Pre-warm SemanticClassifier in the main thread so run_in_executor
+        # workers never race to load multiple SentenceTransformer instances.
+        SemanticClassifier.get_instance()
+        logger.info("Initialization complete — starting task processing.")
+    if _lock is None:
+        _lock = asyncio.Lock()
+
     if not os.path.exists(INPUT_PATH):
         logger.error("Input file not found: %s", INPUT_PATH)
         sys.exit(1)

output/results.json

@@ -1,34 +0,0 @@
-[
-  {
-    "task_id": "factual_1",
-    "answer": "The capital of France is Paris."
-  },
-  {
-    "task_id": "sentiment_1",
-    "answer": "Positive"
-  },
-  {
-    "task_id": "summarization_1",
-    "answer": "Artificial intelligence has significantly impacted various sectors in recent years."
-  },
-  {
-    "task_id": "ner_1",
-    "answer": "[{\"entity\": \"Barack Obama\", \"type\": \"Person\"}, {\"entity\": \"Hawaii\", \"type\": \"Location\"}]"
-  },
-  {
-    "task_id": "math_1",
-    "answer": "4104"
-  },
-  {
-    "task_id": "debug_1",
-    "answer": "def foo(): \n    return 42"
-  },
-  {
-    "task_id": "code_gen_1",
-    "answer": "def add_numbers(num1, num2):\n    return num1 + num2"
-  },
-  {
-    "task_id": "logic_1",
-    "answer": "Yes"
-  }
-]

tests/test_router.py

@@ -113,7 +113,6 @@ async def test_route_api_code() -> None:
         patch("handlers.sentiment.LocalSLMEngine.get_instance"),
         patch("handlers.ner.LocalSLMEngine.get_instance"),
         patch("handlers.summarization.LocalSLMEngine.get_instance"),
-        patch("handlers.code_gen.classify_code_difficulty", return_value="hard"),
         patch("engines.remote_llm.RemoteLLMEngine.generate", new_callable=AsyncMock) as mock_remote,
     ):
         mock_remote.return_value = "def add(a, b): return a + b"