Fix/classifier

Dockerfile

@@ -49,6 +49,6 @@ ENV PYTHONUNBUFFERED=1 \
     LOCAL_MODEL_PATH=/app/models/qwen2.5-3b-instruct-q4_k_m.gguf \
     LOCAL_N_GPU_LAYERS=0 \
     LOCAL_N_THREADS=2 \
-    LOCAL_N_CTX=2048
+    LOCAL_N_CTX=1024
 
 ENTRYPOINT ["/app/entrypoint.sh"]

agent/classifier.py

@@ -6,6 +6,7 @@
 """
 
 import logging
+import re
 from typing import Any
 
 from engines.local_slm import ClassifierEngine
@@ -64,7 +65,7 @@
 def _get_grammar() -> Any:
     global _grammar
     if _grammar is None:
-        from llama_cpp import LlamaGrammar  # type: ignore
+        from llama_cpp import LlamaGrammar  # type: ignore[attr-defined]
 
         _grammar = LlamaGrammar.from_string(_GRAMMAR_STR)
     return _grammar
@@ -73,19 +74,33 @@ def _get_grammar() -> Any:
 def classify(prompt: str) -> str:
     """
     Classify a prompt into one of the routing destinations.
+    Uses regex for Math (safest pattern) to save tokens/latency,
+    and the local LLM for natural language tasks.
     """
     if len(prompt) > _LONG_CONTEXT_THRESHOLD:
         logger.debug("Classifier: long context (%d chars) → %s", len(prompt), ROUTE_API_LONG)
         return ROUTE_API_LONG
 
+    p = prompt.lower()
+
+    # 1. Fast Regex for Math (safest)
+    if re.search(r"\b(calculate|math|arithmetic|percentage|equation|multiply|divide|add|subtract|solve|derivative)\b", p) or re.search(
+        r"\d+\s*[\+\-\*\/]\s*\d+", p
+    ):
+        if not re.search(r"\b(python|script|code|riddle|puzzle|logical|story|function|bug)\b", p):
+            logger.debug("Classifier: matched regex → %s", ROUTE_API_MATH)
+            return ROUTE_API_MATH
+
+    # 2. LLM for the rest
     engine = ClassifierEngine.get_instance()
-    raw = engine.generate(
+    raw_output = engine.generate(
         prompt=f"Task:\n{prompt}\n\nLabel:",
         system_prompt=_SYSTEM_PROMPT,
         max_tokens=8,
         temperature=0.0,
         grammar=_get_grammar(),
-    ).strip()
+    )
+    raw = str(raw_output).strip()
 
     if raw in _LABELS:
         logger.debug("Classifier: LLM → %s", raw)

agent/router.py

@@ -9,6 +9,7 @@
 """
 
 import asyncio
+import concurrent.futures
 import logging
 
 from agent.ast_eval import evaluate_math_expression
@@ -45,8 +46,9 @@
 class AgentRouter:
     """Orchestrates all 4 routing layers for a single prompt."""
 
-    def __init__(self, cache: SemanticCache) -> None:
+    def __init__(self, cache: SemanticCache, executor: concurrent.futures.Executor | None = None) -> None:
         self.cache = cache
+        self.executor = executor
 
         # Instantiate all handlers once (singleton SLM loaded once)
         self._factual = FactualHandler()
@@ -79,12 +81,10 @@ async def route(self, task_id: str, prompt: str) -> str:
             return math_result
 
         # -------------------------------------------------------------------
-        # Layer 2: LLM Classifier (offloaded to thread pool — local llama.cpp
-        # call is synchronous; offloading keeps the event loop responsive so
-        # concurrent API tasks can progress while classification runs).
+        # Layer 2: Hybrid Classifier (Regex for Math, LLM for rest)
         # -------------------------------------------------------------------
         loop = asyncio.get_running_loop()
-        route = await loop.run_in_executor(None, classify, prompt)
+        route = await loop.run_in_executor(self.executor, classify, prompt)
         logger.info("[%s] Route → %s", task_id, route)
 
         # -------------------------------------------------------------------
@@ -104,14 +104,14 @@ async def _dispatch(self, task_id: str, prompt: str, route: str) -> str:
             # tasks waiting for the lock) keep the event loop alive instead of
             # freezing it for the duration of one local generation.
             if route == ROUTE_LOCAL_SENTIMENT:
-                res = await loop.run_in_executor(None, self._sentiment.handle, prompt)
+                res = await loop.run_in_executor(self.executor, self._sentiment.handle, prompt)
                 if res == "__ESCALATE__":
                     logger.info("[%s] Sentiment escalated → remote", task_id)
                     res = await self._remote_general.handle(prompt, category=ROUTE_LOCAL_SENTIMENT)
                 return res
 
             if route == ROUTE_LOCAL_NER:
-                res = await loop.run_in_executor(None, self._ner.handle, prompt)
+                res = await loop.run_in_executor(self.executor, self._ner.handle, prompt)
                 if res == "__ESCALATE__":
                     logger.info("[%s] NER escalated → remote", task_id)
                     res = await self._remote_general.handle(prompt, category=ROUTE_LOCAL_NER)
@@ -120,9 +120,9 @@ async def _dispatch(self, task_id: str, prompt: str, route: str) -> str:
             if route == ROUTE_LOCAL_GENERAL:
                 p = prompt.lower()
                 if any(w in p for w in ["summarize", "summary", "tldr"]):
-                    res = await loop.run_in_executor(None, self._summarization.handle, prompt)
+                    res = await loop.run_in_executor(self.executor, self._summarization.handle, prompt)
                 else:
-                    res = await loop.run_in_executor(None, self._factual.handle, prompt)
+                    res = await loop.run_in_executor(self.executor, self._factual.handle, prompt)
                 if res == "__ESCALATE__":
                     logger.info("[%s] Local general escalated → remote", task_id)
                     res = await self._remote_general.handle(prompt, category=ROUTE_LOCAL_GENERAL)

engines/remote_llm.py

@@ -125,7 +125,7 @@ def __init__(self) -> None:
 
     @retry(
         reraise=True,
-        stop=stop_after_attempt(3),
+        stop=stop_after_attempt(2),
         wait=wait_exponential(multiplier=1, min=2, max=10),
         retry=retry_if_exception_type(
             (
@@ -184,27 +184,20 @@ async def generate(
                 "temperature": temperature,
             }
 
-        timeout = aiohttp.ClientTimeout(total=20, connect=5)
-        # Use a class-level or module-level semaphore to limit concurrency
-        # Initialize it lazily in an async context since asyncio.Semaphore needs an active event loop
-        if self._semaphore is None:
-            self._semaphore = asyncio.Semaphore(15)
-
-        sem: asyncio.Semaphore = self._semaphore
-        async with sem:
-            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()
-                        logger.error("Fireworks API error (status %d): %s", response.status, text)
-                        response.raise_for_status()
-
-                    data = await response.json()
-                    try:
-                        if is_chat:
-                            return str(data["choices"][0]["message"]["content"]).strip()
-                        else:
-                            return str(data["choices"][0]["text"]).strip()
-                    except KeyError as e:
-                        logger.error("Response structure: %s", data)
-                        raise e
+        timeout = aiohttp.ClientTimeout(total=15, 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()
+                    logger.error("Fireworks API error (status %d): %s", response.status, text)
+                    response.raise_for_status()
+
+                data = await response.json()
+                try:
+                    if is_chat:
+                        return str(data["choices"][0]["message"]["content"]).strip()
+                    else:
+                        return str(data["choices"][0]["text"]).strip()
+                except KeyError as e:
+                    logger.error("Response structure: %s", data)
+                    raise e

main.py

@@ -4,6 +4,7 @@
 """
 
 import asyncio
+import concurrent.futures
 import json
 import logging
 import os
@@ -12,7 +13,6 @@
 from dotenv import load_dotenv
 
 from agent.cache import SemanticCache
-from agent.classifier import classify
 from agent.router import AgentRouter
 from agent.schemas import Task
 from agent.watchdog import Watchdog
@@ -41,7 +41,15 @@ def _get_results() -> list[dict[str, str]]:
 
 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)
+    try:
+        answer = await asyncio.wait_for(_router.route(task.task_id, task.prompt), timeout=25.0)
+    except TimeoutError:
+        logger.error("Task %s timed out after 25s", task.task_id)
+        answer = "Error: Execution timed out."
+    except Exception as exc:
+        logger.error("Task %s failed: %s", task.task_id, exc)
+        answer = f"Error: {exc}"
+
     async with _lock:
         _completed.append({"task_id": task.task_id, "answer": answer})
     logger.info("Done: %s", task.task_id)
@@ -50,22 +58,16 @@ async def _process(task: Task) -> None:
 async def main() -> None:
     global _cache, _router, _lock
 
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
+
     # -----------------------------------------------------------------------
     # 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 the LLM classifier with one trivial call. The local model is
-        # already loaded by AgentRouter's handlers; this just primes the
-        # grammar object and the llama.cpp sampler cache so the first real task
-        # does not pay a cold-start tax. Safe to fail — real tasks still route.
-        try:
-            classify("warmup")
-        except Exception as exc:
-            logger.warning("Classifier warmup failed: %s", exc)
+        _router = AgentRouter(cache=_cache, executor=executor)
         logger.info("Initialization complete — starting task processing.")
     if _lock is None:
         _lock = asyncio.Lock()
@@ -114,6 +116,7 @@ async def main() -> None:
         logger.error("Fatal error: %s", exc)
     finally:
         watchdog.stop()
+        executor.shutdown(wait=False)
 
 
 if __name__ == "__main__":