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 1737cf2..bc061d7 100644 --- a/agent/classifier.py +++ b/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) 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 694a1b8..6f9bc38 100644 --- a/engines/remote_llm.py +++ b/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 diff --git a/main.py b/main.py index e1fb43a..ed5fd6d 100644 --- a/main.py +++ b/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,6 +58,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 +67,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 +116,7 @@ async def main() -> None: logger.error("Fatal error: %s", exc) finally: watchdog.stop() + executor.shutdown(wait=False) if __name__ == "__main__":