bstein/titan-iac
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): cap quick chat generation

Browse Source
This commit is contained in:
jenkins 2026-06-29 15:43:32 -03:00
parent f9ea2be3d5
commit d997f281f0
3 changed files with 38 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
services/comms/atlasbot-deployment.yaml
View File

@ -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

51
services/comms/scripts/atlasbot/bot.py
View File

@ -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))

4
services/comms/scripts/tests/test_atlasbot_modes.py
View File

@ -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"