diff --git a/.dockerignore b/.dockerignore index d4d45b4..c9ddf6c 100644 --- a/.dockerignore +++ b/.dockerignore @@ -16,6 +16,10 @@ requirements-dev.txt .idea/ .vscode/ *.swp +.venv/ +venv/ +env/ +.agents/ # Python cache __pycache__/ diff --git a/agent/classifier.py b/agent/classifier.py index 59acd69..d863f61 100644 --- a/agent/classifier.py +++ b/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: diff --git a/agent/router.py b/agent/router.py index 6d8afca..ca19ca8 100644 --- a/agent/router.py +++ b/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) # ------------------------------------------------------------------- diff --git a/engines/remote_llm.py b/engines/remote_llm.py index a57f82f..ec7591b 100644 --- a/engines/remote_llm.py +++ b/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() diff --git a/handlers/code_gen.py b/handlers/code_gen.py index 338017e..0a29872 100644 --- a/handlers/code_gen.py +++ b/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", diff --git a/handlers/debug.py b/handlers/debug.py index ed320d3..a18b583 100644 --- a/handlers/debug.py +++ b/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", diff --git a/main.py b/main.py index 333d703..49e4c9b 100644 --- a/main.py +++ b/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,11 +28,11 @@ 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]]: @@ -39,6 +40,7 @@ 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) async with _lock: _completed.append({"task_id": task.task_id, "answer": answer}) @@ -46,6 +48,23 @@ async def _process(task: Task) -> None: 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) diff --git a/output/results.json b/output/results.json index 105affc..e69de29 100644 --- a/output/results.json +++ b/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" - } -] diff --git a/tests/test_router.py b/tests/test_router.py index 21b6bcd..d3da56a 100644 --- a/tests/test_router.py +++ b/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"