From 33d3989d29f9d49678bc9b92dd70b6b417c44d77 Mon Sep 17 00:00:00 2001 From: "anthropic-code-agent[bot]" <242468646+Claude@users.noreply.github.com> Date: Sun, 26 Jul 2026 13:02:56 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(config):=20resolve=20Ollama?= =?UTF-8?q?=20endpoint=20from=20env=20before=20defaulting=20to=2011434?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Check OLLAMA_BASE_URL, OLLAMA_API_BASE and OLLAMA_HOST before falling back to http://localhost:11434, and reuse the resolved value for the local model, embedder, classifier and attack defaults. Co-authored-by: franconicola <51865029+franconicola@users.noreply.github.com> --- docs/docs/agents/ollama.mdx | 14 ++++ .../attacks/techniques/advprefix/config.py | 5 +- .../indirect_prompt_injection/attack.py | 3 +- .../indirect_prompt_injection/config.py | 3 +- hackagent/attacks/techniques/pair/attack.py | 7 +- hackagent/cli/commands/examples.py | 5 +- hackagent/cli/commands/scan.py | 9 +-- hackagent/config.py | 61 +++++++++++++++- hackagent/router/_chat_registration.py | 5 +- .../router/tracking/category_classifier.py | 3 +- tests/unit/test_config.py | 69 +++++++++++++++++++ 11 files changed, 165 insertions(+), 19 deletions(-) create mode 100644 tests/unit/test_config.py diff --git a/docs/docs/agents/ollama.mdx b/docs/docs/agents/ollama.mdx index 46060436..f4ddcc59 100644 --- a/docs/docs/agents/ollama.mdx +++ b/docs/docs/agents/ollama.mdx @@ -219,6 +219,20 @@ agent = HackAgent( ) ``` +### Non-Default Port + +If your Ollama server does not listen on the default port `11434`, set one of +`OLLAMA_BASE_URL`, `OLLAMA_API_BASE` or `OLLAMA_HOST` (checked in that order). +HackAgent reads them when building the default endpoint for attacker, judge, +classifier and embedder models, so you don't have to override each one: + +```bash +export OLLAMA_HOST="127.0.0.1:11435" +``` + +A value without a scheme defaults to `http://`, and a value without a port +falls back to `11434`. + ## Further Reading - [Ollama Documentation](https://ollama.com/docs) diff --git a/hackagent/attacks/techniques/advprefix/config.py b/hackagent/attacks/techniques/advprefix/config.py index 5f022ae6..71eb1dcb 100644 --- a/hackagent/attacks/techniques/advprefix/config.py +++ b/hackagent/attacks/techniques/advprefix/config.py @@ -21,6 +21,7 @@ DEFAULT_ATTACKER_IDENTIFIER, DEFAULT_FILTER_LEN, DEFAULT_JUDGE_IDENTIFIER, + DEFAULT_LOCAL_MODEL_ENDPOINT, DEFAULT_OUTPUT_DIR, DEFAULT_TIMEOUT, DEFAULT_RUN_ID, @@ -52,7 +53,7 @@ # --- Model Configurations --- "generator": { "identifier": DEFAULT_ATTACKER_IDENTIFIER, - "endpoint": "http://localhost:11434", + "endpoint": DEFAULT_LOCAL_MODEL_ENDPOINT, "system_prompt": DEFAULT_ADVPREFIX_GENERATOR_SYSTEM_PROMPT, "max_tokens": 50, "guided_topk": 50, @@ -61,7 +62,7 @@ "judges": [ { "identifier": DEFAULT_JUDGE_IDENTIFIER, - "endpoint": "http://localhost:11434", + "endpoint": DEFAULT_LOCAL_MODEL_ENDPOINT, "type": "harmbench", } ], diff --git a/hackagent/attacks/techniques/indirect_prompt_injection/attack.py b/hackagent/attacks/techniques/indirect_prompt_injection/attack.py index 7e782886..6b947587 100644 --- a/hackagent/attacks/techniques/indirect_prompt_injection/attack.py +++ b/hackagent/attacks/techniques/indirect_prompt_injection/attack.py @@ -29,6 +29,7 @@ from hackagent.attacks.techniques.base import BaseAttack from hackagent.attacks.shared.router_factory import create_router from hackagent.attacks.shared.response_utils import extract_response_content +from hackagent.config import DEFAULT_EMBEDDER_OPENAI_ENDPOINT from hackagent.router.router import AgentRouter from hackagent.router.tracking.tracker import Tracker from hackagent.server.client import AuthenticatedClient @@ -229,7 +230,7 @@ def get_embeddings( # Fall back to a placeholder so keyless local backends (e.g. Ollama) work: # the OpenAI client requires a non-empty api_key, but local servers ignore it. api_key = config.get("api_key") or os.environ.get("OPENAI_API_KEY") or "not-needed" - raw_endpoint = str(config.get("endpoint", "http://localhost:11434/v1")).strip() + raw_endpoint = str(config.get("endpoint", DEFAULT_EMBEDDER_OPENAI_ENDPOINT)).strip() endpoint = raw_endpoint.rstrip("/") if endpoint.lower().endswith("/embeddings"): # OpenAI client expects API base and appends '/embeddings' internally. diff --git a/hackagent/attacks/techniques/indirect_prompt_injection/config.py b/hackagent/attacks/techniques/indirect_prompt_injection/config.py index f8d201ff..44e1f87f 100644 --- a/hackagent/attacks/techniques/indirect_prompt_injection/config.py +++ b/hackagent/attacks/techniques/indirect_prompt_injection/config.py @@ -10,6 +10,7 @@ from pydantic import Field from hackagent.attacks.techniques.config import ConfigBase, DEFAULT_CONFIG_BASE +from hackagent.config import DEFAULT_EMBEDDER_OPENAI_ENDPOINT DEFAULT_INDIRECT_PROMPT_INJECTION_CONFIG: Dict[str, Any] = { @@ -39,7 +40,7 @@ }, "embedder": { "identifier": "nomic-embed-text", - "endpoint": "http://localhost:11434/v1", + "endpoint": DEFAULT_EMBEDDER_OPENAI_ENDPOINT, "api_key": None, }, }, diff --git a/hackagent/attacks/techniques/pair/attack.py b/hackagent/attacks/techniques/pair/attack.py index 9482f80c..4481bb39 100644 --- a/hackagent/attacks/techniques/pair/attack.py +++ b/hackagent/attacks/techniques/pair/attack.py @@ -28,6 +28,7 @@ from hackagent.attacks.techniques.config import ( DEFAULT_ATTACKER_IDENTIFIER, DEFAULT_JUDGE_IDENTIFIER, + DEFAULT_LOCAL_MODEL_ENDPOINT, ) from hackagent.attacks.objectives import OBJECTIVES from hackagent.attacks.shared.progress import create_progress_bar @@ -285,7 +286,9 @@ def _initialize_attacker_router(self) -> Optional[AgentRouter]: "identifier": attacker_config.get( "identifier", DEFAULT_ATTACKER_IDENTIFIER ), - "endpoint": attacker_config.get("endpoint", "http://localhost:11434"), + "endpoint": attacker_config.get( + "endpoint", DEFAULT_LOCAL_MODEL_ENDPOINT + ), "agent_type": attacker_config.get("agent_type", "OLLAMA"), "thinking": attacker_config.get("thinking"), "max_tokens": attacker_config.get("max_tokens", 500), @@ -334,7 +337,7 @@ def _initialize_judge_router(self) -> Optional[AgentRouter]: router_config = { "identifier": judge_config.get("identifier", DEFAULT_JUDGE_IDENTIFIER), - "endpoint": judge_config.get("endpoint", "http://localhost:11434"), + "endpoint": judge_config.get("endpoint", DEFAULT_LOCAL_MODEL_ENDPOINT), "agent_type": judge_config.get("agent_type", "OLLAMA"), "thinking": judge_config.get("thinking"), "max_tokens": judge_config.get("max_tokens", 4096), diff --git a/hackagent/cli/commands/examples.py b/hackagent/cli/commands/examples.py index 8dc12e48..713e1a10 100644 --- a/hackagent/cli/commands/examples.py +++ b/hackagent/cli/commands/examples.py @@ -27,6 +27,7 @@ from hackagent.cli.config import CLIConfig from hackagent.cli.utils import handle_errors +from hackagent.config import resolve_ollama_base_url console = Console() @@ -309,7 +310,7 @@ def _preflight_ollama_requirements(demo_cfg: dict) -> None: ) raise click.ClickException("Ollama is not installed") - endpoint = demo_cfg.get("agent", {}).get("endpoint") or "http://localhost:11434" + endpoint = demo_cfg.get("agent", {}).get("endpoint") or resolve_ollama_base_url() console.print(f"[cyan]🔎 Checking Ollama server at:[/cyan] {endpoint}") if not _is_ollama_running(str(endpoint)): @@ -504,7 +505,7 @@ def web_example(): "identifier": "gemma3:4b", "type": "harmbench_variant", "agent_type": "ollama", - "endpoint": "http://localhost:11434", + "endpoint": resolve_ollama_base_url(), } ], } diff --git a/hackagent/cli/commands/scan.py b/hackagent/cli/commands/scan.py index 1251dcdf..89e728f0 100644 --- a/hackagent/cli/commands/scan.py +++ b/hackagent/cli/commands/scan.py @@ -18,7 +18,6 @@ from __future__ import annotations -import os import time from typing import Any, Dict, List, Optional, Tuple @@ -36,6 +35,7 @@ handle_errors, load_config_file, ) +from hackagent.config import resolve_ollama_base_url from hackagent.router.discovery.scanner import ( DEFAULT_PLANNER_MODEL, PlannerError, @@ -88,9 +88,10 @@ def _provider_endpoint(model: str) -> str: """Return the api_base URL for a LiteLLM ``model`` id (by provider prefix).""" m = (model or "").strip() prefix = m.split("/", 1)[0].lower() if "/" in m else "" - if prefix in ("ollama", "ollama_chat") or not prefix: - return os.environ.get("OLLAMA_API_BASE") or "http://localhost:11434" - return _PROVIDER_ENDPOINTS.get(prefix, "http://localhost:11434") + endpoint = _PROVIDER_ENDPOINTS.get(prefix) + if prefix in ("ollama", "ollama_chat") or endpoint is None: + return resolve_ollama_base_url() + return endpoint def _extract_asr(results: Any) -> Optional[float]: diff --git a/hackagent/config.py b/hackagent/config.py index 8451a640..8eee7ef9 100644 --- a/hackagent/config.py +++ b/hackagent/config.py @@ -17,10 +17,54 @@ from __future__ import annotations +import os + # --------------------------------------------------------------------------- # Local Ollama defaults (no API key required) # --------------------------------------------------------------------------- +# Environment variables consulted (in order) to locate the local Ollama server +# before falling back to the upstream default host/port. ``OLLAMA_HOST`` is the +# variable Ollama itself honours, so a user who moved the server off 11434 gets +# picked up automatically. +OLLAMA_BASE_URL_ENV_VARS = ("OLLAMA_BASE_URL", "OLLAMA_API_BASE", "OLLAMA_HOST") + +DEFAULT_OLLAMA_HOST = "localhost" +DEFAULT_OLLAMA_PORT = "11434" + + +def _normalize_ollama_base_url(raw: str) -> str: + """Normalise an ``OLLAMA_HOST``-style value into a full base URL. + + Accepts ``http://host:port``, ``host:port``, ``host`` and ``:port`` forms. + A missing scheme defaults to ``http``, a missing host to ``localhost`` and a + missing port to Ollama's default port (``https`` values keep the implicit + 443 instead). + """ + value = raw.strip().rstrip("/") + scheme, sep, remainder = value.partition("://") + if not sep: + scheme, remainder = "http", value + remainder = remainder.lstrip("/") + authority, slash, path = remainder.partition("/") + if authority.startswith(":"): + authority = f"{DEFAULT_OLLAMA_HOST}{authority}" + # Only append the default port for a bare http host (no port, not IPv6). + # https values are left alone so they keep the implicit 443. + if scheme == "http" and ":" not in authority and not authority.endswith("]"): + authority = f"{authority}:{DEFAULT_OLLAMA_PORT}" + return f"{scheme}://{authority}{slash}{path}" + + +def resolve_ollama_base_url() -> str: + """Return the local Ollama base URL, honouring environment overrides.""" + for env_var in OLLAMA_BASE_URL_ENV_VARS: + raw = os.environ.get(env_var, "").strip() + if raw: + return _normalize_ollama_base_url(raw) + return f"http://{DEFAULT_OLLAMA_HOST}:{DEFAULT_OLLAMA_PORT}" + + # Local Ollama default model. Uncensored so it won't refuse to generate # red-team prompts. Pull: ``ollama pull huihui_ai/gemma-4-abliterated:12b``. DEFAULT_LOCAL_MODEL = "huihui_ai/gemma-4-abliterated:12b" @@ -32,11 +76,11 @@ # Default local embedder served by Ollama (used by any attack that needs an # embedder, e.g. the RAG Attack and AutoDAN-Turbo strategy retrieval). DEFAULT_EMBEDDER_IDENTIFIER = "embeddinggemma" -DEFAULT_EMBEDDER_ENDPOINT = "http://localhost:11434" +DEFAULT_EMBEDDER_ENDPOINT = resolve_ollama_base_url() DEFAULT_EMBEDDER_AGENT_TYPE = "OLLAMA" # OpenAI-compatible base URL exposed by Ollama (used by the RAG Attack, which # embeds through an OpenAI-compatible client and posts to ``/v1/embeddings``). -DEFAULT_EMBEDDER_OPENAI_ENDPOINT = "http://localhost:11434/v1" +DEFAULT_EMBEDDER_OPENAI_ENDPOINT = f"{resolve_ollama_base_url()}/v1" # Ollama ignores the key but the OpenAI client requires a non-empty value. DEFAULT_EMBEDDER_OPENAI_API_KEY = "ollama" @@ -45,7 +89,7 @@ # form; callers that split identifier/endpoint/agent_type want DEFAULT_LOCAL_MODEL. DEFAULT_LOCAL_LITELLM_MODEL = f"{OLLAMA_PROVIDER_PREFIX}/{DEFAULT_LOCAL_MODEL}" -DEFAULT_LOCAL_MODEL_ENDPOINT = "http://localhost:11434" +DEFAULT_LOCAL_MODEL_ENDPOINT = resolve_ollama_base_url() DEFAULT_LOCAL_AGENT_TYPE = "OLLAMA" # Local role identifiers — attacker / judge / category-classifier all default @@ -69,12 +113,23 @@ DEFAULT_REMOTE_JUDGE_IDENTIFIER = "hackagent-judge" __all__ = [ + # ollama endpoint resolution + "OLLAMA_BASE_URL_ENV_VARS", + "DEFAULT_OLLAMA_HOST", + "DEFAULT_OLLAMA_PORT", + "resolve_ollama_base_url", # local model "DEFAULT_LOCAL_MODEL", "OLLAMA_PROVIDER_PREFIX", "DEFAULT_LOCAL_LITELLM_MODEL", "DEFAULT_LOCAL_MODEL_ENDPOINT", "DEFAULT_LOCAL_AGENT_TYPE", + # local embedder + "DEFAULT_EMBEDDER_IDENTIFIER", + "DEFAULT_EMBEDDER_ENDPOINT", + "DEFAULT_EMBEDDER_AGENT_TYPE", + "DEFAULT_EMBEDDER_OPENAI_ENDPOINT", + "DEFAULT_EMBEDDER_OPENAI_API_KEY", # local roles "DEFAULT_ATTACKER_IDENTIFIER", "DEFAULT_JUDGE_IDENTIFIER", diff --git a/hackagent/router/_chat_registration.py b/hackagent/router/_chat_registration.py index 3decae10..315ced71 100644 --- a/hackagent/router/_chat_registration.py +++ b/hackagent/router/_chat_registration.py @@ -23,6 +23,7 @@ import os from typing import Any, Dict, Optional +from hackagent.config import resolve_ollama_base_url from hackagent.logger import get_logger from hackagent.router import envelope as _envelope from hackagent.router.provider_config import ProviderConfig @@ -35,12 +36,10 @@ # These helpers cover the small adapter-class quirks that used to live # in ``OpenAIAgent.__init__`` and ``OllamaAgent.__init__``. -_OLLAMA_DEFAULT_ENDPOINT = "http://localhost:11434" - def _normalise_ollama_endpoint(endpoint: Optional[str]) -> str: """Resolve & normalise the Ollama endpoint URL the way OllamaAgent did.""" - resolved = endpoint or os.environ.get("OLLAMA_BASE_URL", _OLLAMA_DEFAULT_ENDPOINT) + resolved = endpoint or resolve_ollama_base_url() resolved = resolved.rstrip("/") for suffix in ("/api/generate", "/api/chat", "/api/tags", "/api/show", "/api"): if resolved.endswith(suffix): diff --git a/hackagent/router/tracking/category_classifier.py b/hackagent/router/tracking/category_classifier.py index 6f361b15..18722f98 100644 --- a/hackagent/router/tracking/category_classifier.py +++ b/hackagent/router/tracking/category_classifier.py @@ -286,11 +286,12 @@ def _resolve_config(config: Optional[Dict[str, Any]]) -> Dict[str, Any]: # Imported lazily to avoid a router↔attacks import cycle at load time. from hackagent.attacks.techniques.config import ( DEFAULT_CATEGORY_CLASSIFIER_IDENTIFIER, + DEFAULT_CATEGORY_CLASSIFIER_ENDPOINT, ) resolved: Dict[str, Any] = { "identifier": DEFAULT_CATEGORY_CLASSIFIER_IDENTIFIER, - "endpoint": "http://localhost:11434", + "endpoint": DEFAULT_CATEGORY_CLASSIFIER_ENDPOINT, "agent_type": "OLLAMA", "api_key": None, "max_tokens": 100, diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 00000000..b0e005dc --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,69 @@ +# Copyright 2026 - AI4I. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for ``hackagent/config.py`` Ollama endpoint resolution.""" + +import os +import unittest +from unittest.mock import patch + +from hackagent.config import resolve_ollama_base_url + + +class TestResolveOllamaBaseUrl(unittest.TestCase): + def _resolve(self, env): + with patch.dict(os.environ, env, clear=True): + return resolve_ollama_base_url() + + def test_defaults_to_localhost_11434(self): + self.assertEqual(self._resolve({}), "http://localhost:11434") + + def test_ollama_base_url_wins(self): + self.assertEqual( + self._resolve({"OLLAMA_BASE_URL": "http://ollama:11435"}), + "http://ollama:11435", + ) + + def test_ollama_api_base_is_used(self): + self.assertEqual( + self._resolve({"OLLAMA_API_BASE": "http://localhost:11500"}), + "http://localhost:11500", + ) + + def test_ollama_host_without_scheme(self): + self.assertEqual( + self._resolve({"OLLAMA_HOST": "127.0.0.1:11435"}), + "http://127.0.0.1:11435", + ) + + def test_ollama_host_port_only(self): + self.assertEqual( + self._resolve({"OLLAMA_HOST": ":11435"}), "http://localhost:11435" + ) + + def test_ollama_host_without_port_gets_default_port(self): + self.assertEqual( + self._resolve({"OLLAMA_HOST": "my-ollama"}), "http://my-ollama:11434" + ) + + def test_trailing_slash_is_stripped(self): + self.assertEqual( + self._resolve({"OLLAMA_BASE_URL": "http://localhost:11435/"}), + "http://localhost:11435", + ) + + def test_blank_env_var_falls_through(self): + self.assertEqual( + self._resolve({"OLLAMA_BASE_URL": " ", "OLLAMA_HOST": "host:11499"}), + "http://host:11499", + ) + + def test_https_scheme_keeps_implicit_port(self): + self.assertEqual( + self._resolve({"OLLAMA_HOST": "https://remote.example.com"}), + "https://remote.example.com", + ) + + +if __name__ == "__main__": + unittest.main()