RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: Add Ollama Cloud integration with ollama_cloud/ prefix (#372) 

* feat: Add Ollama Cloud integration with ollama_cloud/ prefix

- Added support for Ollama Cloud models via AsyncOpenAI
- Models use ollama_cloud/ prefix (e.g., ollama_cloud/gpt-oss:120b)
- Reads OLLAMA_API_KEY and OLLAMA_API_BASE for cloud authentication
- Added predefined Ollama Cloud models to /model-show
- Filtered obsolete ollama/*-cloud models from LiteLLM database
- Updated authentication headers in completer, banner, and toolbar
- Added concise English documentation in docs/providers/
- Adapted to new global model cache architecture from #371

Modified files:
- src/cai/sdk/agents/models/openai_chatcompletions.py
- src/cai/repl/commands/model.py (adapted to global cache)
- src/cai/util.py
- src/cai/repl/commands/completer.py
- src/cai/repl/ui/banner.py
- src/cai/repl/ui/toolbar.py
- docs/providers/ollama.md
- docs/providers/ollama_cloud.md

All existing functionality preserved (backward compatible).

* fix: Add missing Ollama Cloud models and get_ollama_auth_headers

- Added Ollama Cloud models to get_predefined_model_categories()
- Added get_ollama_auth_headers() function to util.py
- Added Ollama Cloud to category_to_provider mapping
- Imported get_ollama_auth_headers in model.py

This fixes the issue where predefined models weren't showing in /model-show.

* fix: Display predefined models first in /model-show

- Added loop to display predefined models (Alias, Claude, OpenAI, DeepSeek, Ollama Cloud) before LiteLLM models
- Models #1-14 now correctly show predefined models
- LiteLLM models start from #15+ as expected
- Added get_ollama_auth_headers() function in util.py for Ollama Cloud auth
- Fixed model numbering to be consistent with global cache

This resolves the issue where predefined models were skipped and only LiteLLM models appeared starting from #15.

* fix: Restore correct import of cai.caibench instead of pentestperf

During rebase, the import was incorrectly changed from 'import cai.caibench as ptt'
to 'import pentestperf as ptt'. This commit restores the correct import.

The original code uses cai.caibench, not an external pentestperf module.
This commit is contained in:
pzabalegui 2025-12-10 15:42:48 +01:00 committed by GitHub
parent 7c2267d2b9
commit da1c6202e5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 329 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
docs/providers/ollama.md
View File

@ -1,5 +1,7 @@
# Ollama Configuration
## Ollama Local (Self-hosted)
#### [Ollama Integration](https://ollama.com/)
For local models using Ollama, add the following to your .env:
@ -9,3 +11,28 @@ OLLAMA_API_BASE=http://localhost:8000/v1 # note, maybe you have a different endp
```
Make sure that the Ollama server is running and accessible at the specified base URL. You can swap the model with any other supported by your local Ollama instance.
## Ollama Cloud
For cloud models using Ollama Cloud (no GPU required), add the following to your .env:
```bash
# API Key from ollama.com
OLLAMA_API_KEY=your_api_key_here
OLLAMA_API_BASE=https://ollama.com
# Cloud model (note the ollama_cloud/ prefix)
CAI_MODEL=ollama_cloud/gpt-oss:120b
```
**Requirements:**
1. Create an account at [ollama.com](https://ollama.com)
2. Generate an API key from your profile
3. Use models with `ollama_cloud/` prefix (e.g., `ollama_cloud/gpt-oss:120b`)
**Key differences:**
- Prefix: `ollama_cloud/` (cloud) vs `ollama/` (local)
- API Key: Required for cloud, not needed for local
- Endpoint: `https://ollama.com/v1` (cloud) vs `http://localhost:8000/v1` (local)
See [Ollama Cloud documentation](ollama_cloud.md) for detailed setup instructions.

79
docs/providers/ollama_cloud.md Normal file
View File

@ -0,0 +1,79 @@
# Ollama Cloud
Run large language models without local GPU using Ollama's cloud service.
## Quick Start
### 1. Get API Key
- Create account at [ollama.com](https://ollama.com)
- Generate API key from your profile
### 2. Configure `.env`
```bash
OLLAMA_API_KEY=your_api_key_here
OLLAMA_API_BASE=https://ollama.com
CAI_MODEL=ollama_cloud/gpt-oss:120b
```
### 3. Run
```bash
cai
```
## Available Models
View in CAI with `/model-show` under "Ollama Cloud" category:
- `ollama_cloud/gpt-oss:120b` - General purpose 120B model
- `ollama_cloud/llama3.3:70b` - Llama 3.3 70B
- `ollama_cloud/qwen2.5:72b` - Qwen 2.5 72B
- `ollama_cloud/deepseek-v3:671b` - DeepSeek V3 671B
More models at [ollama.com/library](https://ollama.com/library).
## Model Selection
```bash
# By name
CAI> /model ollama_cloud/gpt-oss:120b
# By number (after /model-show)
CAI> /model 3
```
## Local vs Cloud
| Feature | Local | Cloud |
|---------|-------|-------|
| Prefix | `ollama/` | `ollama_cloud/` |
| API Key | Not required | Required |
| Endpoint | `http://localhost:8000/v1` | `https://ollama.com/v1` |
| GPU | Required | Not required |
## Troubleshooting
**Unauthorized error**: Verify `OLLAMA_API_KEY` is set correctly
**Path not found**: Ensure `OLLAMA_API_BASE=https://ollama.com` (without `/v1`)
**Model not listed**: Check model prefix is `ollama_cloud/`, not `ollama/`
## Validation
Test connection with curl:
```bash
curl https://ollama.com/v1/chat/completions \
-H "Authorization: Bearer $OLLAMA_API_KEY" \
-H "Content-Type: application/json" \
-d '{"model": "gpt-oss:120b", "messages": [{"role": "user", "content": "test"}]}'
```
## References
- [Ollama Cloud Docs](https://ollama.com/docs/cloud)
- [Model Library](https://ollama.com/library)
- [Get API Key](https://ollama.com/settings/keys)

12
src/cai/repl/commands/completer.py
View File

@ -184,8 +184,18 @@ class FuzzyCommandCompleter(Completer):
try:
# Get Ollama models with a short timeout to prevent hanging
api_base = get_ollama_api_base()
# Add authentication headers for Ollama Cloud if using OPENAI_BASE_URL
headers = {}
if "ollama.com" in api_base:
api_key = os.getenv("OPENAI_API_KEY")
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
response = requests.get(
f"{api_base.replace('/v1', '')}/api/tags", timeout=0.5)
f"{api_base.replace('/v1', '')}/api/tags",
headers=headers,
timeout=0.5)
if response.status_code == 200:
data = response.json()

90
src/cai/repl/commands/model.py
View File

@ -12,7 +12,7 @@ import requests # pylint: disable=import-error
from rich.console import Console # pylint: disable=import-error
from rich.table import Table # pylint: disable=import-error
from rich.panel import Panel # pylint: disable=import-error
from cai.util import get_ollama_api_base, COST_TRACKER
from cai.util import get_ollama_api_base, get_ollama_auth_headers, COST_TRACKER
from cai.repl.commands.base import Command, register_command
console = Console()
@ -99,6 +99,32 @@ def get_predefined_model_categories() -> Dict[str, List[Dict[str, str]]]:
"name": "deepseek-r1",
"description": "DeepSeek's specialized reasoning model"
}
],
"Ollama Cloud": [
{
"name": "ollama_cloud/gpt-oss:120b",
"description": (
"Ollama Cloud - Large 120B parameter model (no GPU required)"
)
},
{
"name": "ollama_cloud/llama3.3:70b",
"description": (
"Ollama Cloud - Llama 3.3 70B model (no GPU required)"
)
},
{
"name": "ollama_cloud/qwen2.5:72b",
"description": (
"Ollama Cloud - Qwen 2.5 72B model (no GPU required)"
)
},
{
"name": "ollama_cloud/deepseek-v3:671b",
"description": (
"Ollama Cloud - DeepSeek V3 671B model (no GPU required)"
)
}
]
}
@ -117,7 +143,8 @@ def get_all_predefined_models() -> List[Dict[str, Any]]:
"Alias": "OpenAI", # Alias models use OpenAI as base
"Anthropic Claude": "Anthropic",
"OpenAI": "OpenAI",
"DeepSeek": "DeepSeek"
"DeepSeek": "DeepSeek",
"Ollama Cloud": "Ollama Cloud"
}
for category, models in model_categories.items():
@ -175,7 +202,11 @@ def load_all_available_models() -> tuple[List[str], List[Dict[str, Any]]]:
try:
response = requests.get(LITELLM_URL, timeout=5)
if response.status_code == 200:
litellm_names = sorted(response.json().keys())
# Filter out obsolete Ollama Cloud models (replaced by ollama_cloud/ prefix)
litellm_names = [
model_name for model_name in sorted(response.json().keys())
if not (model_name.startswith("ollama/") and "-cloud" in model_name)
]
except Exception: # pylint: disable=broad-except
pass
@ -184,7 +215,17 @@ def load_all_available_models() -> tuple[List[str], List[Dict[str, Any]]]:
ollama_names = []
try:
api_base = get_ollama_api_base()
response = requests.get(f"{api_base.replace('/v1', '')}/api/tags", timeout=1)
ollama_base = api_base.replace('/v1', '')
# Add authentication headers for Ollama Cloud if needed
headers = {}
is_cloud = "ollama.com" in api_base
timeout = 5 if is_cloud else 1 # Cloud needs more time
if is_cloud:
headers = get_ollama_auth_headers()
response = requests.get(f"{ollama_base}/api/tags", headers=headers, timeout=timeout)
if response.status_code == 200:
data = response.json()
ollama_data = data.get('models', data.get('items', []))
@ -499,7 +540,46 @@ class ModelShowCommand(Command):
total_models = 0
displayed_models = 0
# Process and display models (use global cache for numbering)
# First, add predefined models (Alias, Claude, OpenAI, DeepSeek, Ollama Cloud)
predefined_models = get_all_predefined_models()
for model in predefined_models:
model_name = model["name"]
# Skip if search term provided and not in model name
if search_term and search_term not in model_name.lower():
continue
displayed_models += 1
total_models += 1
# Find index from global cache
try:
model_index = _GLOBAL_MODEL_CACHE.index(model_name) + 1
except ValueError:
continue
# Format pricing info
input_cost_str = (
f"${model['input_cost']:.2f}"
if model['input_cost'] is not None else "Unknown"
)
output_cost_str = (
f"${model['output_cost']:.2f}"
if model['output_cost'] is not None else "Unknown"
)
# Add row to table
model_table.add_row(
str(model_index),
model_name,
model["provider"],
"N/A", # max_tokens
input_cost_str,
output_cost_str,
model.get("description", "")
)
# Process and display LiteLLM models (use global cache for numbering)
for model_name, model_info in sorted(model_data.items()):
# Find the model index from global cache
try:

15
src/cai/repl/ui/banner.py
View File

@ -76,12 +76,19 @@ def get_supported_models_count():
# Try to get Ollama models count
try:
ollama_api_base = os.getenv(
"OLLAMA_API_BASE",
"http://host.docker.internal:8000/v1"
)
from cai.util import get_ollama_api_base
ollama_api_base = get_ollama_api_base()
# Add authentication headers for Ollama Cloud if using OPENAI_BASE_URL
headers = {}
if "ollama.com" in ollama_api_base:
api_key = os.getenv("OPENAI_API_KEY")
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
ollama_response = requests.get(
f"{ollama_api_base.replace('/v1', '')}/api/tags",
headers=headers,
timeout=1
)

17
src/cai/repl/ui/toolbar.py
View File

@ -96,11 +96,20 @@ def update_toolbar_in_background():
ollama_status = "unavailable"
try:
# Get Ollama models with a short timeout to prevent hanging
api_base = os.getenv(
"OLLAMA_API_BASE",
"http://host.docker.internal:8000/v1")
from cai.util import get_ollama_api_base
api_base = get_ollama_api_base()
# Add authentication headers for Ollama Cloud if using OPENAI_BASE_URL
headers = {}
if "ollama.com" in api_base:
api_key = os.getenv("OPENAI_API_KEY")
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
response = requests.get(
f"{api_base.replace('/v1', '')}/api/tags", timeout=0.5)
f"{api_base.replace('/v1', '')}/api/tags",
headers=headers,
timeout=0.5)
if response.status_code == 200:
data = response.json()

63
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -2708,7 +2708,23 @@ class OpenAIChatCompletionsModel(Model):
provider = model_str.split("/")[0]
# Apply provider-specific configurations
if provider == "deepseek":
if provider == "ollama_cloud":
# Ollama Cloud configuration
ollama_api_key = os.getenv("OLLAMA_API_KEY")
ollama_api_base = os.getenv("OLLAMA_API_BASE", "https://ollama.com")
if ollama_api_key:
kwargs["api_key"] = ollama_api_key
if ollama_api_base:
kwargs["api_base"] = ollama_api_base
# Drop params not supported by Ollama
litellm.drop_params = True
kwargs.pop("parallel_tool_calls", None)
kwargs.pop("store", None)
if not converted_tools:
kwargs.pop("tool_choice", None)
elif provider == "deepseek":
litellm.drop_params = True
kwargs.pop("parallel_tool_calls", None)
kwargs.pop("store", None) # DeepSeek doesn't support store parameter
@ -2846,6 +2862,51 @@ class OpenAIChatCompletionsModel(Model):
max_retries = 3
retry_count = 0
# Check if this is Ollama Cloud (ollama_cloud/ prefix)
# Ollama Cloud is OpenAI-compatible, so we bypass LiteLLM to avoid parsing issues
is_ollama_cloud = "ollama_cloud/" in model_str
if is_ollama_cloud:
# Use AsyncOpenAI client directly for Ollama Cloud
# Ollama Cloud is fully OpenAI-compatible at /v1/chat/completions
try:
# Configure the client with Ollama Cloud settings
ollama_api_key = os.getenv("OLLAMA_API_KEY") or os.getenv("OPENAI_API_KEY")
ollama_base_url = os.getenv("OLLAMA_API_BASE", "https://ollama.com")
# Ensure the URL has /v1 for OpenAI compatibility
if not ollama_base_url.endswith("/v1"):
ollama_base_url = f"{ollama_base_url}/v1"
# Create a temporary client configured for Ollama Cloud
ollama_client = AsyncOpenAI(
api_key=ollama_api_key,
base_url=ollama_base_url
)
# Remove the ollama_cloud/ prefix from the model name
clean_model = kwargs["model"].replace("ollama_cloud/", "")
kwargs["model"] = clean_model
# Remove LiteLLM-specific parameters
kwargs.pop("extra_headers", None)
kwargs.pop("api_key", None)
kwargs.pop("api_base", None)
kwargs.pop("custom_llm_provider", None)
# Call Ollama Cloud using OpenAI-compatible API
if stream:
return await ollama_client.chat.completions.create(**kwargs)
else:
return await ollama_client.chat.completions.create(**kwargs)
except Exception as e:
# If Ollama Cloud fails, raise with helpful message
raise Exception(
f"Error connecting to Ollama Cloud: {str(e)}\n"
f"Verify OLLAMA_API_KEY and OLLAMA_API_BASE are configured correctly."
) from e
while retry_count < max_retries:
try:
if self.is_ollama:

44
src/cai/util.py
View File

@ -31,8 +31,14 @@ from wasabi import color
from cai import is_pentestperf_available
# Import caibench (pentestperf) if available
if is_pentestperf_available():
import pentestperf as ptt
import cai.caibench as ptt
PTT_AVAILABLE = True
else:
ptt = None
PTT_AVAILABLE = False
import signal
# Global timing variables for tracking active and idle time
@ -744,8 +750,36 @@ install_pretty()
def get_ollama_api_base():
"""Get the Ollama API base URL from environment variable or default to localhost:8000."""
return os.environ.get("OLLAMA_API_BASE", "http://localhost:8000/v1")
"""Get the Ollama API base URL from environment variable or default to localhost:8000.
Supports both:
- OLLAMA_API_BASE: For local Ollama instances (e.g., http://localhost:8000/v1)
- OPENAI_BASE_URL: For Ollama Cloud or other OpenAI-compatible services (e.g., https://ollama.com/api/v1)
"""
# First check OLLAMA_API_BASE for local Ollama
ollama_base = os.environ.get("OLLAMA_API_BASE")
if ollama_base:
return ollama_base
# Then check OPENAI_BASE_URL for Ollama Cloud or other services
openai_base = os.environ.get("OPENAI_BASE_URL")
if openai_base and "ollama.com" in openai_base:
return openai_base
# Default to local Ollama
return "http://localhost:8000/v1"
def get_ollama_auth_headers():
"""Get authentication headers for Ollama Cloud if API key is set.
Returns:
Dictionary with Authorization header if API key exists, empty dict otherwise
"""
api_key = os.getenv("OLLAMA_API_KEY") or os.getenv("OPENAI_API_KEY")
if api_key:
return {"Authorization": f"Bearer {api_key}"}
return {}
def load_prompt_template(template_path):
@ -4340,6 +4374,10 @@ def setup_ctf():
print(color("CTF name not provided, necessary to run CTF", fg="white", bg="red"))
sys.exit(1)
if not PTT_AVAILABLE or ptt is None:
print(color("pentestperf module not available, cannot setup CTF", fg="white", bg="red"))
sys.exit(1)
print(
color("Setting up CTF: ", fg="black", bg="yellow")
+ color(ctf_name, fg="black", bg="yellow")