diff --git a/atlasbot/config.py b/atlasbot/config.py index 1b0316a..f74dec5 100644 --- a/atlasbot/config.py +++ b/atlasbot/config.py @@ -61,6 +61,9 @@ class Settings: nats_stream: str nats_subject: str nats_result_bucket: str + nats_timeout_fast_sec: int + nats_timeout_smart_sec: int + nats_timeout_genius_sec: int fast_max_angles: int smart_max_angles: int @@ -72,6 +75,7 @@ class Settings: smart_llm_calls_max: int genius_llm_calls_max: int llm_limit_multiplier: float + debug_pipeline: bool @dataclass(frozen=True) @@ -154,6 +158,9 @@ def load_settings() -> Settings: nats_stream=os.getenv("ATLASBOT_NATS_STREAM", "atlasbot"), nats_subject=os.getenv("ATLASBOT_NATS_SUBJECT", "atlasbot.requests"), nats_result_bucket=os.getenv("ATLASBOT_NATS_RESULTS", "atlasbot_results"), + nats_timeout_fast_sec=_env_int("ATLASBOT_NATS_TIMEOUT_FAST_SEC", "180"), + nats_timeout_smart_sec=_env_int("ATLASBOT_NATS_TIMEOUT_SMART_SEC", "480"), + nats_timeout_genius_sec=_env_int("ATLASBOT_NATS_TIMEOUT_GENIUS_SEC", "900"), fast_max_angles=_env_int("ATLASBOT_FAST_MAX_ANGLES", "2"), smart_max_angles=_env_int("ATLASBOT_SMART_MAX_ANGLES", "5"), genius_max_angles=_env_int("ATLASBOT_GENIUS_MAX_ANGLES", "9"), @@ -164,4 +171,5 @@ def load_settings() -> Settings: smart_llm_calls_max=_env_int("ATLASBOT_SMART_LLM_CALLS_MAX", "17"), genius_llm_calls_max=_env_int("ATLASBOT_GENIUS_LLM_CALLS_MAX", "32"), llm_limit_multiplier=_env_float("ATLASBOT_LLM_LIMIT_MULTIPLIER", "1.5"), + debug_pipeline=_env_bool("ATLASBOT_DEBUG_PIPELINE", "false"), ) diff --git a/atlasbot/queue/nats.py b/atlasbot/queue/nats.py index a274659..418a8b2 100644 --- a/atlasbot/queue/nats.py +++ b/atlasbot/queue/nats.py @@ -43,7 +43,14 @@ class QueueManager: sub = await self._nc.subscribe(reply) envelope = {"reply": reply, "payload": payload} await self._js.publish(self._settings.nats_subject, json.dumps(envelope).encode()) - msg = await sub.next_msg(timeout=300) + mode = (payload.get("mode") or "").strip().lower() + if mode == "genius": + timeout = self._settings.nats_timeout_genius_sec + elif mode == "smart": + timeout = self._settings.nats_timeout_smart_sec + else: + timeout = self._settings.nats_timeout_fast_sec + msg = await sub.next_msg(timeout=timeout) await sub.unsubscribe() return json.loads(msg.data.decode())