mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-15 17:25:26 -04:00
feat(ai): add OpenRouter and Ollama Cloud providers (#231)
Co-authored-by: Alex Kenley <Alex.Kenley@threatvectorsecurity.com>
This commit is contained in:
@@ -11,15 +11,35 @@ def normalize_base(url: str) -> str:
|
||||
for suffix in ["/models", "/chat/completions", "/completions", "/v1/messages"]:
|
||||
if url.endswith(suffix):
|
||||
url = url[: -len(suffix)].rstrip("/")
|
||||
for suffix in ["/chat", "/tags", "/generate"]:
|
||||
if url.endswith("/api" + suffix):
|
||||
url = url[: -len(suffix)].rstrip("/")
|
||||
return url
|
||||
|
||||
|
||||
def _detect_provider(url: str) -> str:
|
||||
parsed = urlparse(url or "")
|
||||
host = parsed.hostname or ""
|
||||
path = (parsed.path or "").rstrip("/")
|
||||
if host.endswith("ollama.com") or (parsed.port == 11434 and (path == "/api" or path.startswith("/api/"))):
|
||||
return "ollama"
|
||||
if "anthropic.com" in (url or ""):
|
||||
return "anthropic"
|
||||
return "openai"
|
||||
|
||||
|
||||
def _ollama_api_root(base: str) -> str:
|
||||
base = (base or "").strip().rstrip("/")
|
||||
parsed = urlparse(base)
|
||||
host = parsed.hostname or ""
|
||||
path = (parsed.path or "").rstrip("/")
|
||||
if path.endswith("/api"):
|
||||
return base
|
||||
if host.endswith("ollama.com"):
|
||||
return f"{parsed.scheme}://{parsed.netloc}/api"
|
||||
return base
|
||||
|
||||
|
||||
def build_chat_url(base: str) -> str:
|
||||
provider = _detect_provider(base)
|
||||
if provider == "anthropic":
|
||||
@@ -27,9 +47,18 @@ def build_chat_url(base: str) -> str:
|
||||
if host.endswith("anthropic.com") and base.rstrip("/").endswith("/v1"):
|
||||
base = base.rstrip("/")[:-3].rstrip("/")
|
||||
return base + "/v1/messages"
|
||||
if provider == "ollama":
|
||||
return _ollama_api_root(base) + "/chat"
|
||||
return base + "/chat/completions"
|
||||
|
||||
|
||||
def build_models_url(base: str) -> str:
|
||||
provider = _detect_provider(base)
|
||||
if provider == "ollama":
|
||||
return _ollama_api_root(base) + "/tags"
|
||||
return base + "/models"
|
||||
|
||||
|
||||
def build_headers(api_key, base: str) -> dict:
|
||||
if not api_key:
|
||||
return {}
|
||||
@@ -52,6 +81,9 @@ class TestNormalizeBase:
|
||||
def test_strips_v1_messages(self):
|
||||
assert normalize_base("https://api.anthropic.com/v1/messages") == "https://api.anthropic.com"
|
||||
|
||||
def test_strips_ollama_native_chat(self):
|
||||
assert normalize_base("https://ollama.com/api/chat") == "https://ollama.com/api"
|
||||
|
||||
def test_trailing_slash(self):
|
||||
assert normalize_base("https://api.openai.com/v1/") == "https://api.openai.com/v1"
|
||||
|
||||
@@ -78,6 +110,20 @@ class TestBuildChatUrl:
|
||||
def test_local_endpoint(self):
|
||||
assert build_chat_url("http://localhost:8000/v1") == "http://localhost:8000/v1/chat/completions"
|
||||
|
||||
def test_ollama_cloud_native_api(self):
|
||||
assert build_chat_url("https://ollama.com/api") == "https://ollama.com/api/chat"
|
||||
|
||||
def test_ollama_cloud_root_adds_api(self):
|
||||
assert build_chat_url("https://ollama.com") == "https://ollama.com/api/chat"
|
||||
|
||||
|
||||
class TestBuildModelsUrl:
|
||||
def test_openai_models(self):
|
||||
assert build_models_url("https://api.openai.com/v1") == "https://api.openai.com/v1/models"
|
||||
|
||||
def test_ollama_tags(self):
|
||||
assert build_models_url("https://ollama.com/api") == "https://ollama.com/api/tags"
|
||||
|
||||
|
||||
class TestBuildHeaders:
|
||||
def test_no_key(self):
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Regression tests for native Ollama Cloud provider handling."""
|
||||
import httpx
|
||||
|
||||
from src import llm_core
|
||||
|
||||
|
||||
def test_detects_ollama_cloud_native_provider():
|
||||
assert llm_core._detect_provider("https://ollama.com/api") == "ollama"
|
||||
assert llm_core._detect_provider("https://ollama.com/api/chat") == "ollama"
|
||||
|
||||
|
||||
def test_llm_call_posts_native_ollama_payload(monkeypatch):
|
||||
seen = {}
|
||||
|
||||
def fake_post(url, headers=None, json=None, timeout=None):
|
||||
seen["url"] = url
|
||||
seen["headers"] = headers
|
||||
seen["json"] = json
|
||||
seen["timeout"] = timeout
|
||||
request = httpx.Request("POST", url)
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
json={"message": {"content": "OK"}, "done": True},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(llm_core.httpx, "post", fake_post)
|
||||
|
||||
result = llm_core.llm_call(
|
||||
"https://ollama.com/api",
|
||||
"gpt-oss:120b-test",
|
||||
[{"role": "user", "content": "Say OK"}],
|
||||
temperature=0.2,
|
||||
max_tokens=7,
|
||||
headers={"Authorization": "Bearer ollama-key"},
|
||||
timeout=11,
|
||||
)
|
||||
|
||||
assert result == "OK"
|
||||
assert seen["url"] == "https://ollama.com/api/chat"
|
||||
assert seen["headers"]["Authorization"] == "Bearer ollama-key"
|
||||
assert seen["json"]["stream"] is False
|
||||
assert seen["json"]["options"] == {"temperature": 0.2, "num_predict": 7}
|
||||
@@ -65,6 +65,9 @@ class TestMatchProviderCurated:
|
||||
def test_xai_url(self):
|
||||
assert _match_provider_curated("https://api.x.ai/v1", "openai") == "xai"
|
||||
|
||||
def test_ollama_url(self):
|
||||
assert _match_provider_curated("https://ollama.com/api", "openai") == "ollama"
|
||||
|
||||
def test_no_url_match_returns_provider(self):
|
||||
assert _match_provider_curated("https://localhost:1234", "openai") == "openai"
|
||||
|
||||
@@ -263,6 +266,26 @@ class TestSetupProbeSafety:
|
||||
assert _probe_endpoint("https://api.anthropic.com/v1", "good-key") == ["claude-sonnet-4-5"]
|
||||
assert seen == ["https://api.anthropic.com/v1/models"]
|
||||
|
||||
def test_ollama_cloud_probe_uses_native_tags_endpoint(self, monkeypatch):
|
||||
monkeypatch.setattr(endpoint_resolver, "resolve_url", lambda url: url, raising=False)
|
||||
monkeypatch.setattr(model_routes, "_normalize_base", lambda url: url.rstrip("/"))
|
||||
seen = []
|
||||
|
||||
def fake_get(url, headers=None, timeout=None):
|
||||
seen.append((url, headers))
|
||||
request = httpx.Request("GET", url)
|
||||
response = httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
json={"models": [{"name": "gpt-oss:120b"}, {"model": "qwen3:235b"}]},
|
||||
)
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(model_routes.httpx, "get", fake_get)
|
||||
|
||||
assert _probe_endpoint("https://ollama.com/api", "ollama-key") == ["gpt-oss:120b", "qwen3:235b"]
|
||||
assert seen == [("https://ollama.com/api/tags", {"Authorization": "Bearer ollama-key"})]
|
||||
|
||||
def test_unkeyed_anthropic_probe_can_use_curated_fallback(self, monkeypatch):
|
||||
monkeypatch.setattr(endpoint_resolver, "resolve_url", lambda url: url, raising=False)
|
||||
monkeypatch.setattr(model_routes, "_normalize_base", lambda url: url.rstrip("/"))
|
||||
|
||||
Reference in New Issue
Block a user