From 59faedd0a24e49ac5b08eefa88c0a6303d7e490c Mon Sep 17 00:00:00 2001 From: Yoshio Date: Mon, 13 Jul 2026 16:39:01 +0700 Subject: [PATCH 1/2] fix(timeout): Implement hybrid regex/LLM classifier and thread pool isolation --- Dockerfile | 2 +- agent/classifier.py | 14 ++++++++++++-- agent/router.py | 18 +++++++++--------- engines/remote_llm.py | 4 ++-- main.py | 24 ++++++++++++++---------- 5 files changed, 38 insertions(+), 24 deletions(-) diff --git a/Dockerfile b/Dockerfile index 47552e6..5887415 100644 --- a/Dockerfile +++ b/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"] diff --git a/agent/classifier.py b/agent/classifier.py index 75dc308..c00e1ff 100644 --- a/agent/classifier.py +++ b/agent/classifier.py @@ -7,6 +7,7 @@ import logging import os +import re from engines.local_slm import LocalSLMEngine @@ -60,12 +61,10 @@ _grammar = None - def _get_grammar(): global _grammar if _grammar is None: from llama_cpp import LlamaGrammar - _grammar = LlamaGrammar.from_string(_GRAMMAR_STR) return _grammar @@ -73,11 +72,22 @@ def _get_grammar(): 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 = LocalSLMEngine.get_instance() raw = engine.generate( prompt=f"Task:\n{prompt}\n\nLabel:", diff --git a/agent/router.py b/agent/router.py index 828b9b4..7729773 100644 --- a/agent/router.py +++ b/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) diff --git a/engines/remote_llm.py b/engines/remote_llm.py index ec7591b..fc60192 100644 --- a/engines/remote_llm.py +++ b/engines/remote_llm.py @@ -124,7 +124,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( ( @@ -183,7 +183,7 @@ async def generate( "temperature": temperature, } - timeout = aiohttp.ClientTimeout(total=20, connect=5) + 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: diff --git a/main.py b/main.py index e1fb43a..2410604 100644 --- a/main.py +++ b/main.py @@ -4,6 +4,7 @@ """ import asyncio +import concurrent.futures import json import logging import os @@ -41,7 +42,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 asyncio.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,6 +59,8 @@ 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(). @@ -57,15 +68,7 @@ async def main() -> None: 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 +117,7 @@ async def main() -> None: logger.error("Fatal error: %s", exc) finally: watchdog.stop() + executor.shutdown(wait=False) if __name__ == "__main__": From 377b941393c92bea77b003dc0a9710f30f6b3c3a Mon Sep 17 00:00:00 2001 From: Yoshio Date: Mon, 13 Jul 2026 16:49:59 +0700 Subject: [PATCH 2/2] fix(optimize): remove _call_lock & add threadpool --- agent/classifier.py | 16 ++++++++++------ main.py | 5 ++--- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/agent/classifier.py b/agent/classifier.py index c00e1ff..975bd47 100644 --- a/agent/classifier.py +++ b/agent/classifier.py @@ -6,8 +6,8 @@ """ import logging -import os import re +from typing import Any from engines.local_slm import LocalSLMEngine @@ -61,10 +61,12 @@ _grammar = None -def _get_grammar(): + +def _get_grammar() -> Any: global _grammar if _grammar is None: - from llama_cpp import LlamaGrammar + from llama_cpp import LlamaGrammar # type: ignore[attr-defined] + _grammar = LlamaGrammar.from_string(_GRAMMAR_STR) return _grammar @@ -80,10 +82,12 @@ def classify(prompt: str) -> str: 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): + 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 diff --git a/main.py b/main.py index 2410604..ed5fd6d 100644 --- a/main.py +++ b/main.py @@ -13,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 @@ -44,13 +43,13 @@ async def _process(task: Task) -> None: assert _router is not None and _lock is not None try: answer = await asyncio.wait_for(_router.route(task.task_id, task.prompt), timeout=25.0) - except asyncio.TimeoutError: + 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)