diff --git a/services/comms/atlasbot-deployment.yaml b/services/comms/atlasbot-deployment.yaml index 7432ce145..58761aa5a 100644 --- a/services/comms/atlasbot-deployment.yaml +++ b/services/comms/atlasbot-deployment.yaml @@ -131,6 +131,8 @@ spec: value: "600" - name: ATLASBOT_QUICK_TIME_BUDGET_SEC value: "15" + - name: ATLASBOT_FAST_NUM_PREDICT + value: "48" - name: ATLASBOT_SMART_TIME_BUDGET_SEC value: "45" - name: ATLASBOT_GENIUS_TIME_BUDGET_SEC diff --git a/services/comms/scripts/atlasbot/bot.py b/services/comms/scripts/atlasbot/bot.py index ef82775d7..749cc126b 100644 --- a/services/comms/scripts/atlasbot/bot.py +++ b/services/comms/scripts/atlasbot/bot.py @@ -54,6 +54,7 @@ THINKING_INTERVAL_SEC = int(os.environ.get("ATLASBOT_THINKING_INTERVAL_SEC", "12 QUICK_TIME_BUDGET_SEC = float(os.environ.get("ATLASBOT_QUICK_TIME_BUDGET_SEC", "15")) SMART_TIME_BUDGET_SEC = float(os.environ.get("ATLASBOT_SMART_TIME_BUDGET_SEC", "45")) GENIUS_TIME_BUDGET_SEC = float(os.environ.get("ATLASBOT_GENIUS_TIME_BUDGET_SEC", "180")) +FAST_NUM_PREDICT = int(os.environ.get("ATLASBOT_FAST_NUM_PREDICT", "48")) OLLAMA_RETRIES = int(os.environ.get("ATLASBOT_OLLAMA_RETRIES", "2")) OLLAMA_SERIALIZE = os.environ.get("ATLASBOT_OLLAMA_SERIALIZE", "true").lower() != "false" @@ -492,6 +493,12 @@ def _mode_heartbeat_sec(mode: str) -> int: return max(5, min(THINKING_INTERVAL_SEC, int(max(5.0, budget / 3.0)))) +def _mode_ollama_options(mode: str) -> dict[str, Any] | None: + if _normalize_mode(mode) == "fast": + return {"num_predict": max(8, FAST_NUM_PREDICT)} + return None + + # Matrix HTTP helper. def req(method: str, path: str, token: str | None = None, body=None, timeout=60, base: str | None = None): url = (base or BASE) + path @@ -3085,17 +3092,19 @@ def _ollama_call_safe( system_override: str | None = None, model: str | None = None, timeout: float | None = None, + options: dict[str, Any] | None = None, ) -> str: try: - return _ollama_call( - hist_key, - prompt, - context=context, - use_history=False, - system_override=system_override, - model=model, - timeout=timeout, - ) + kwargs: dict[str, Any] = { + "context": context, + "use_history": False, + "system_override": system_override, + "model": model, + "timeout": timeout, + } + if options: + kwargs["options"] = options + return _ollama_call(hist_key, prompt, **kwargs) except Exception: return fallback @@ -4321,6 +4330,7 @@ def _open_ended_fast_single( system_override=_open_ended_system(), model=model, timeout=_mode_ollama_timeout_sec("fast"), + options=_mode_ollama_options("fast"), ) if not _has_body_lines(reply): reply = _ollama_call_safe( @@ -4331,6 +4341,7 @@ def _open_ended_fast_single( system_override=_open_ended_system(), model=model, timeout=_mode_ollama_timeout_sec("fast"), + options=_mode_ollama_options("fast"), ) fallback = _fallback_fact_answer(prompt, fallback_context or context) if fallback and (_is_quantitative_prompt(prompt) or not _has_body_lines(reply)): @@ -4512,6 +4523,7 @@ def _non_cluster_reply(prompt: str, *, history_lines: list[str], mode: str) -> s system_override=system, model=model, timeout=_mode_ollama_timeout_sec(mode), + options=_mode_ollama_options(mode), ) reply = re.sub(r"\bconfidence\s*:\s*(high|medium|low)\b\.?\s*", "", reply, flags=re.IGNORECASE).strip() return _ensure_scores(reply) @@ -4856,6 +4868,7 @@ def _ollama_call( system_override: str | None = None, model: str | None = None, timeout: float | None = None, + options: dict[str, Any] | None = None, ) -> str: system = system_override or ( "System: You are Atlas, the Titan lab assistant for Atlas/Othrys. " @@ -4891,6 +4904,8 @@ def _ollama_call( model_name = model or MODEL request_timeout = timeout if timeout is not None else OLLAMA_TIMEOUT_SEC payload = {"model": model_name, "messages": messages, "stream": False} + if options: + payload["options"] = options headers = {"Content-Type": "application/json"} if API_KEY: headers["x-api-key"] = API_KEY @@ -4932,18 +4947,20 @@ def ollama_reply( use_history: bool = True, model: str | None = None, timeout: float | None = None, + options: dict[str, Any] | None = None, ) -> str: last_error = None for attempt in range(max(1, OLLAMA_RETRIES + 1)): try: - return _ollama_call( - hist_key, - prompt, - context=context, - use_history=use_history, - model=model, - timeout=timeout, - ) + kwargs: dict[str, Any] = { + "context": context, + "use_history": use_history, + "model": model, + "timeout": timeout, + } + if options: + kwargs["options"] = options + return _ollama_call(hist_key, prompt, **kwargs) except Exception as exc: # noqa: BLE001 last_error = exc time.sleep(min(4, 2 ** attempt)) diff --git a/services/comms/scripts/tests/test_atlasbot_modes.py b/services/comms/scripts/tests/test_atlasbot_modes.py index 15f7d5367..aa54b0b43 100644 --- a/services/comms/scripts/tests/test_atlasbot_modes.py +++ b/services/comms/scripts/tests/test_atlasbot_modes.py @@ -82,7 +82,7 @@ class AtlasbotModeTests(TestCase): ] captured: dict[str, object] = {} - def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None): + def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None, options=None): captured["model"] = model captured["timeout"] = timeout captured["context"] = context @@ -115,7 +115,7 @@ class AtlasbotModeTests(TestCase): ] seen: list[tuple[str, float]] = [] - def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None): + def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None, options=None): seen.append((str(model), float(timeout or 0))) return "Atlas has a clear standout because the worker spread is healthy. Confidence: high"