Enhance support for Ollama models and update environment configuration

- Updated `.env.example`, `CONTRIBUTING.md`, and `README.md` to include instructions for using Ollama models with optional local API base configuration.
- Modified CLI error handling to provide specific messages based on the selected model provider.
- Enhanced `has_supported_provider_key` function to accommodate Ollama, allowing it to run without cloud API keys.
- Added new test cases to verify the correct behavior of provider key requirements and model-specific configurations.

Author: Av7danger

.env.example

@@ -4,3 +4,6 @@ SECNODE_LLM=openai/gpt-4o
 # Provider API keys
 OPENAI_API_KEY=
 ANTHROPIC_API_KEY=
+
+# Optional local endpoint when using ollama/* models
+OLLAMA_API_BASE=http://localhost:11434

CONTRIBUTING.md

@@ -41,8 +41,14 @@ Thank you for your interest in contributing to SecNode API! This guide will help
    # Or Anthropic
    export SECNODE_LLM="anthropic/claude-3-5-sonnet-20241022"
    export ANTHROPIC_API_KEY="your-anthropic-key"
+
+   # Or Ollama
+   export SECNODE_LLM="ollama/llama3.1"
+   export OLLAMA_API_BASE="http://localhost:11434" # optional
    ```
-   Provider keys are model-specific. See [LiteLLM provider docs](https://docs.litellm.ai/docs/providers) for other providers.
+   Provider credentials are model-specific. `openai/*` requires `OPENAI_API_KEY`,
+   `anthropic/*` requires `ANTHROPIC_API_KEY`, and `ollama/*` can run locally without
+   cloud API keys. See [LiteLLM provider docs](https://docs.litellm.ai/docs/providers) for other providers.
 
 5. **Run SecNode in development mode**
    ```bash

README.md

@@ -77,7 +77,16 @@ export SECNODE_LLM="anthropic/claude-3-5-sonnet-20241022"
 export ANTHROPIC_API_KEY="your-anthropic-key"
 ```
 
-Provider keys are model-specific. See [LiteLLM providers](https://docs.litellm.ai/docs/providers).
+Or with Ollama:
+
+```bash
+export SECNODE_LLM="ollama/llama3.1"
+export OLLAMA_API_BASE="http://localhost:11434" # optional, defaults to localhost if omitted
+```
+
+Provider credentials are model-specific. `openai/*` requires `OPENAI_API_KEY`, `anthropic/*`
+requires `ANTHROPIC_API_KEY`, and `ollama/*` can run locally without cloud API keys.
+See [LiteLLM providers](https://docs.litellm.ai/docs/providers).
 
 ## Quick Start
 

src/secnodeapi/ai/llm_client.py

@@ -16,19 +16,30 @@
 async def call_llm(system_prompt: str, user_prompt: str, temperature: float = 0.2) -> str:
     """Call LiteLLM with retries and basic rate-limit backoff."""
     model = os.getenv("SECNODE_LLM", DEFAULT_MODEL)
+    provider = model.split("/", 1)[0].lower() if "/" in model else model.lower()
     max_retries = 5
     base_delay = 2.0
 
     for attempt in range(max_retries):
         try:
-            response = await acompletion(
-                model=model,
-                messages=[
+            completion_kwargs = {
+                "model": model,
+                "messages": [
                     {"role": "system", "content": system_prompt},
                     {"role": "user", "content": user_prompt},
                 ],
-                temperature=temperature,
-                response_format={"type": "json_object"},
+                "temperature": temperature,
+            }
+            if provider != "ollama":
+                completion_kwargs["response_format"] = {"type": "json_object"}
+
+            if provider == "ollama":
+                api_base = os.getenv("OLLAMA_API_BASE", "").strip()
+                if api_base:
+                    completion_kwargs["api_base"] = api_base
+
+            response = await acompletion(
+                **completion_kwargs
             )
             return response.choices[0].message.content
         except Exception as e:

src/secnodeapi/cli.py

@@ -3,6 +3,7 @@
 import argparse
 import asyncio
 import json
+import os
 import warnings
 from pathlib import Path
 from typing import List
@@ -153,9 +154,27 @@ def _write_dry_run_output(tests: list, dry_run_output: str) -> None:
 def _require_provider_key(schema_only: bool) -> None:
     if schema_only or has_supported_provider_key():
         return
+    model = os.getenv("SECNODE_LLM", "openai/gpt-4o").strip().lower()
+    provider = model.split("/", 1)[0] if "/" in model else model
+
+    if provider == "openai":
+        message = "Provider API key required. Set OPENAI_API_KEY for SECNODE_LLM openai/* models."
+    elif provider == "anthropic":
+        message = (
+            "Provider API key required. Set ANTHROPIC_API_KEY for SECNODE_LLM anthropic/* models."
+        )
+    elif provider == "ollama":
+        message = (
+            "OLLAMA provider selected but unavailable configuration detected. "
+            "Set SECNODE_LLM=ollama/<model> and optionally OLLAMA_API_BASE."
+        )
+    else:
+        message = (
+            "Provider API key required. Set provider-specific credentials to match SECNODE_LLM."
+        )
+
     logger.error(
-        "Provider API key required. Set OPENAI_API_KEY or ANTHROPIC_API_KEY "
-        "to match SECNODE_LLM."
+        message
     )
     raise SystemExit(1)
 

src/secnodeapi/config.py

@@ -1,8 +1,8 @@
 """
 Runtime configuration for SecNode CLI execution.
 """
-from dataclasses import dataclass
 import os
+from dataclasses import dataclass
 
 
 @dataclass(frozen=True)
@@ -17,5 +17,15 @@ def verify_ssl(self) -> bool:
 
 
 def has_supported_provider_key() -> bool:
-    """Return True when a supported provider API key is configured."""
+    """Return True when provider requirements are met for the selected model."""
+    model = os.getenv("SECNODE_LLM", "openai/gpt-4o").strip().lower()
+    provider = model.split("/", 1)[0] if "/" in model else model
+
+    if provider == "openai":
+        return bool(os.getenv("OPENAI_API_KEY"))
+    if provider == "anthropic":
+        return bool(os.getenv("ANTHROPIC_API_KEY"))
+    if provider == "ollama":
+        return True
+
     return bool(os.getenv("OPENAI_API_KEY") or os.getenv("ANTHROPIC_API_KEY"))

tests/test_cli_and_config.py

@@ -15,12 +15,36 @@ def test_runtime_config_verify_ssl_defaults() -> None:
 def test_has_supported_provider_key(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
     monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    monkeypatch.delenv("SECNODE_LLM", raising=False)
     assert has_supported_provider_key() is False
 
     monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
     assert has_supported_provider_key() is True
 
 
+def test_has_supported_provider_key_respects_selected_provider(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+
+    monkeypatch.setenv("SECNODE_LLM", "anthropic/claude-3-5-sonnet-20241022")
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    assert has_supported_provider_key() is False
+
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-test")
+    assert has_supported_provider_key() is True
+
+
+def test_has_supported_provider_key_allows_ollama_without_cloud_keys(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("SECNODE_LLM", "ollama/llama3.1")
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    assert has_supported_provider_key() is True
+
+
 def test_parse_auth_header_only() -> None:
     headers = cli.parse_auth("Authorization: Bearer token", None)
     assert headers == {"Authorization": "Bearer token"}
@@ -55,6 +79,7 @@ def test_parse_identities_file(tmp_path: Path) -> None:
 def test_require_provider_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
     monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    monkeypatch.setenv("SECNODE_LLM", "openai/gpt-4o")
 
     with pytest.raises(SystemExit):
         cli._require_provider_key(schema_only=False)
@@ -68,6 +93,15 @@ def test_require_provider_key_allows_schema_only(
     cli._require_provider_key(schema_only=True)
 
 
+def test_require_provider_key_allows_ollama(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    monkeypatch.setenv("SECNODE_LLM", "ollama/llama3.1")
+    cli._require_provider_key(schema_only=False)
+
+
 def test_parse_args_with_dry_run_output(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(
         "sys.argv",

tests/test_cli_main_and_llm.py

@@ -88,13 +88,60 @@ async def test_call_llm_success(monkeypatch) -> None:
     message = types.SimpleNamespace(content='{"ok":true}')
     choice = types.SimpleNamespace(message=message)
     response = types.SimpleNamespace(choices=[choice])
+    captured = {}
 
     async def fake_completion(**kwargs):
+        captured.update(kwargs)
         return response
 
+    monkeypatch.delenv("SECNODE_LLM", raising=False)
+    monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
     monkeypatch.setattr("secnodeapi.ai.llm_client.acompletion", fake_completion)
     out = await llm_client.call_llm("sys", "user")
     assert out == '{"ok":true}'
+    assert captured["model"] == "openai/gpt-4o"
+    assert captured["response_format"] == {"type": "json_object"}
+
+
+@pytest.mark.asyncio
+async def test_call_llm_ollama_uses_optional_api_base(monkeypatch) -> None:
+    message = types.SimpleNamespace(content='{"ok":true}')
+    choice = types.SimpleNamespace(message=message)
+    response = types.SimpleNamespace(choices=[choice])
+    captured = {}
+
+    async def fake_completion(**kwargs):
+        captured.update(kwargs)
+        return response
+
+    monkeypatch.setenv("SECNODE_LLM", "ollama/llama3.1")
+    monkeypatch.setenv("OLLAMA_API_BASE", "http://localhost:11434")
+    monkeypatch.setattr("secnodeapi.ai.llm_client.acompletion", fake_completion)
+    out = await llm_client.call_llm("sys", "user")
+    assert out == '{"ok":true}'
+    assert captured["model"] == "ollama/llama3.1"
+    assert captured["api_base"] == "http://localhost:11434"
+    assert "response_format" not in captured
+
+
+@pytest.mark.asyncio
+async def test_call_llm_ollama_without_api_base(monkeypatch) -> None:
+    message = types.SimpleNamespace(content='{"ok":true}')
+    choice = types.SimpleNamespace(message=message)
+    response = types.SimpleNamespace(choices=[choice])
+    captured = {}
+
+    async def fake_completion(**kwargs):
+        captured.update(kwargs)
+        return response
+
+    monkeypatch.setenv("SECNODE_LLM", "ollama/llama3.1")
+    monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
+    monkeypatch.setattr("secnodeapi.ai.llm_client.acompletion", fake_completion)
+    out = await llm_client.call_llm("sys", "user")
+    assert out == '{"ok":true}'
+    assert "api_base" not in captured
+    assert "response_format" not in captured
 
 
 @pytest.mark.asyncio