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

Merge branch 'main' into main

This commit is contained in:
UnaiAlias 2025-12-19 12:54:12 +01:00 committed by GitHub
parent 608762ff21 d086cc4871
commit 03459d9c20
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
19 changed files with 621 additions and 229 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
docs/cai_architecture.md
View File

@ -87,7 +87,7 @@ For more details, including examples and implementation guidance, see the [Agent
You may find different [tools](src/cai/tools). They are grouped in 6 major categories inspired by the security kill chain[2]:
1. Reconnaissance and weaponization - *reconnaissance* (crypto, listing, etc,)
1. Reconnaissance and weaponization - *reconnaissance* (crypto, listing, etc.)
2. Exploitation - *exploitation*
3. Privilege escalation - *escalation*
4. Lateral movement - *lateral*
@ -118,14 +118,14 @@ wherein:
When building `Patterns`, we generally classify them among one of the following categories, though others exist:
| **Agentic** `Pattern` **categories** | **Description** |
|--------------------|------------------------|
| `Swarm` (Decentralized) | Agents share tasks and self-assign responsibilities without a central orchestrator. Handoffs occur dynamically. *An example of a peer-to-peer agentic pattern is the `CTF Agentic Pattern`, which involves a team of agents working together to solve a CTF challenge with dynamic handoffs.* |
| `Hierarchical` | A top-level agent (e.g., "PlannerAgent") assigns tasks via structured handoffs to specialized sub-agents. Alternatively, the structure of the agents is harcoded into the agentic pattern with pre-defined handoffs. |
| **Agentic** `Pattern` **categories** | **Description** |
|--------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `Swarm` (Decentralized) | Agents share tasks and self-assign responsibilities without a central orchestrator. Handoffs occur dynamically. *An example of a peer-to-peer agentic pattern is the `CTF Agentic Pattern`, which involves a team of agents working together to solve a CTF challenge with dynamic handoffs.* |
| `Hierarchical` | A top-level agent (e.g., "PlannerAgent") assigns tasks via structured handoffs to specialized sub-agents. Alternatively, the structure of the agents is hardcoded into the agentic pattern with pre-defined handoffs. |
| `Chain-of-Thought` (Sequential Workflow) | A structured pipeline where Agent A produces an output, hands it to Agent B for reuse or refinement, and so on. Handoffs follow a linear sequence. *An example of a chain-of-thought agentic pattern is the `ReasonerAgent`, which involves a Reasoning-type LLM that provides context to the main agent to solve a CTF challenge with a linear sequence.*[1] |
| `Auction-Based` (Competitive Allocation) | Agents "bid" on tasks based on priority, capability, or cost. A decision agent evaluates bids and hands off tasks to the best-fit agent. |
| `Recursive` | A single agent continuously refines its own output, treating itself as both executor and evaluator, with handoffs (internal or external) to itself. *An example of a recursive agentic pattern is the `CodeAgent` (when used as a recursive agent), which continuously refines its own output by executing code and updating its own instructions.* |
| `Parallelization` | Multiple agents run in parallel, each handling different subtasks or independent inputs simultaneously. This approach speeds up processing when tasks do not depend on each other. *For example, you can launch several agents to analyze different log files or scan multiple IP addresses at the same time, leveraging concurrency to improve efficiency.* |
| `Auction-Based` (Competitive Allocation) | Agents "bid" on tasks based on priority, capability, or cost. A decision agent evaluates bids and hands off tasks to the best-fit agent. |
| `Recursive` | A single agent continuously refines its own output, treating itself as both executor and evaluator, with handoffs (internal or external) to itself. *An example of a recursive agentic pattern is the `CodeAgent` (when used as a recursive agent), which continuously refines its own output by executing code and updating its own instructions.* |
| `Parallelization` | Multiple agents run in parallel, each handling different subtasks or independent inputs simultaneously. This approach speeds up processing when tasks do not depend on each other. *For example, you can launch several agents to analyze different log files or scan multiple IP addresses at the same time, leveraging concurrency to improve efficiency.* |
Moreover in this new version we could orchestrate agents and add decision mechanism in several ways. See [Orchestrating multiple agents](multi_agent.md)

2
docs/cai_pro_contact.md
View File

@ -336,7 +336,7 @@ Review these resources to answer common questions:
[CAI Terms & Conditions](https://aliasrobotics.com/terms-and-conditions.php)
**Privacy Policy:**
GDPR compliant - data processed in EU only
GDPR-compliant - data processed in EU only
**License Information:**
- CAI FREE: [Open source license](https://github.com/aliasrobotics/cai/blob/main/LICENSE)

2
docs/cli/commands_reference.md
View File

@ -1216,7 +1216,7 @@ CAI> /agent bug_bounter_agent ; test https://target.com ; /cost
---
### Auto-loading Queue from File
### Autoloading Queue from File
Load and execute prompts automatically on startup.

2
docs/context.md
View File

@ -11,7 +11,7 @@ This is represented via the [`RunContextWrapper`][cai.sdk.agents.run_context.Run
1. You create any Python object you want. A common pattern is to use a dataclass or a Pydantic object.
2. You pass that object to the various run methods (e.g. `Runner.run(..., **context=whatever**))`).
3. All your tool calls, lifecycle hooks etc will be passed a wrapper object, `RunContextWrapper[T]`, where `T` represents your context object type which you can access via `wrapper.context`.
3. All your tool calls, lifecycle hooks, etc. will be passed a wrapper object, `RunContextWrapper[T]`, where `T` represents your context object type which you can access via `wrapper.context`.
The **most important** thing to be aware of: every agent, tool function, lifecycle, etc for a given agent run must use the same _type_ of context.

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)

2
pyproject.toml
View File

@ -1,6 +1,6 @@
[project]
name = "cai-framework"
version = "0.5.9"
version = "0.5.10"
description = "Cybersecurity AI Framework"
readme = "README.md"
requires-python = ">=3.9"

2
src/cai/agents/__init__.py
View File

@ -36,7 +36,7 @@ where:
| **Agentic Pattern** | **Description** |
|--------------------|------------------------|
| `Swarm` (Decentralized) | Agents share tasks and self-assign responsibilities without a central orchestrator. Handoffs occur dynamically. *An example of a peer-to-peer agentic pattern is the `CTF Agentic Pattern`, which involves a team of agents working together to solve a CTF challenge with dynamic handoffs.* |
| `Hierarchical` | A top-level agent (e.g., "PlannerAgent") assigns tasks via structured handoffs to specialized sub-agents. Alternatively, the structure of the agents is harcoded into the agentic pattern with pre-defined handoffs. |
| `Hierarchical` | A top-level agent (e.g., "PlannerAgent") assigns tasks via structured handoffs to specialized sub-agents. Alternatively, the structure of the agents is hardcoded into the agentic pattern with pre-defined handoffs. |
| `Chain-of-Thought` (Sequential Workflow) | A structured pipeline where Agent A produces an output, hands it to Agent B for reuse or refinement, and so on. Handoffs follow a linear sequence. *An example of a chain-of-thought agentic pattern is the `ReasonerAgent`, which involves a Reasoning-type LLM that provides context to the main agent to solve a CTF challenge with a linear sequence.*[^1] |
| `Auction-Based` (Competitive Allocation) | Agents "bid" on tasks based on priority, capability, or cost. A decision agent evaluates bids and hands off tasks to the best-fit agent. |
| `Recursive` | A single agent continuously refines its own output, treating itself as both executor and evaluator, with handoffs (internal or external) to itself. *An example of a recursive agentic pattern is the `CodeAgent` (when used as a recursive agent), which continuously refines its own output by executing code and updating its own instructions.* |

2
src/cai/agents/web_pentester.py
View File

@ -13,6 +13,7 @@ from cai.agents.guardrails import get_security_guardrails
# Core tools
from cai.tools.reconnaissance.generic_linux_command import generic_linux_command
from cai.tools.reconnaissance.exec_code import execute_code
from cai.tools.web.headers import web_request_framework
# Optional OSINT search (requires PERPLEXITY_API_KEY)
from cai.tools.web.search_web import make_web_search_with_explanation
@ -27,6 +28,7 @@ web_pentester_system_prompt = load_prompt_template("prompts/system_web_pentester
tools = [
generic_linux_command, # shell one-liners (httpie/curl/waybackurls/etc if installed)
execute_code, # light parsing/diffing/payload crafting
web_request_framework, # HTTP request + response/header security analysis
]
# Conditional: add web search helper when available

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

363
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()
@ -22,6 +22,10 @@ LITELLM_URL = (
"model_prices_and_context_window.json"
)
# Global cache shared between /model and /model-show commands
_GLOBAL_MODEL_CACHE = []
_GLOBAL_MODEL_NUMBERS = {}
def get_predefined_model_categories() -> Dict[str, List[Dict[str, str]]]:
"""Get the predefined model categories as the single source of truth.
@ -95,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)"
)
}
]
}
@ -113,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():
@ -155,6 +186,57 @@ def get_predefined_model_names() -> List[str]:
return [model["name"] for model in get_all_predefined_models()]
def load_all_available_models() -> tuple[List[str], List[Dict[str, Any]]]:
"""Load all available models (predefined + LiteLLM + Ollama) in consistent order.
This ensures /model and /model-show use the same numbering.
Returns:
Tuple of (all_model_names, ollama_models_data)
"""
# Predefined models
predefined = [model["name"] for model in get_all_predefined_models()]
# LiteLLM models
litellm_names = []
try:
response = requests.get(LITELLM_URL, timeout=5)
if response.status_code == 200:
# 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
# Ollama models
ollama_data = []
ollama_names = []
try:
api_base = get_ollama_api_base()
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', []))
ollama_names = [m.get('name', '') for m in ollama_data if m.get('name')]
except Exception: # pylint: disable=broad-except
pass
all_models = predefined + litellm_names + ollama_names
return all_models, ollama_data
class ModelCommand(Command):
"""Command for viewing and changing the current LLM model."""
@ -196,29 +278,21 @@ class ModelCommand(Command):
Returns:
bool: True if the model was changed successfully
"""
# Get all predefined models from the shared source of truth
# pylint: disable=invalid-name
ALL_MODELS = get_all_predefined_models()
# Also fetch LiteLLM model names to make numbering consistent with /model-show
litellm_model_names = []
try:
response = requests.get(LITELLM_URL, timeout=5)
if response.status_code == 200:
litellm_data = response.json()
# Add LiteLLM model names (sorted for consistency with /model-show)
litellm_model_names = sorted(litellm_data.keys())
except Exception: # pylint: disable=broad-except
# Silently fail if LiteLLM is not available
pass
# Update cached models to include all models for number selection (consistent with /model-show)
predefined_model_names = [model["name"] for model in ALL_MODELS]
self.cached_models = predefined_model_names + litellm_model_names
self.cached_model_numbers = {
# Load all models with consistent numbering and update global cache
global _GLOBAL_MODEL_CACHE, _GLOBAL_MODEL_NUMBERS
_GLOBAL_MODEL_CACHE, ollama_models_data = load_all_available_models()
_GLOBAL_MODEL_NUMBERS = {
str(i): model_name
for i, model_name in enumerate(self.cached_models, 1)
for i, model_name in enumerate(_GLOBAL_MODEL_CACHE, 1)
}
self.cached_models = _GLOBAL_MODEL_CACHE
self.cached_model_numbers = _GLOBAL_MODEL_NUMBERS
# Get predefined and litellm counts for display
ALL_MODELS = get_all_predefined_models()
predefined_model_names = [model["name"] for model in ALL_MODELS]
litellm_model_names = [m for m in self.cached_models[len(predefined_model_names):]
if m not in [d.get('name') for d in ollama_models_data]]
if not args: # pylint: disable=too-many-nested-blocks
# Display current model
@ -272,61 +346,35 @@ class ModelCommand(Command):
model["description"]
)
# Ollama models (if available)
# pylint: disable=too-many-nested-blocks
try:
# Get Ollama models with a short timeout to prevent hanging
api_base = get_ollama_api_base()
ollama_base = api_base.replace('/v1', '')
response = requests.get(
f"{ollama_base}/api/tags",
timeout=1
)
if response.status_code == 200:
data = response.json()
ollama_models = []
if 'models' in data:
ollama_models = data['models']
else:
# Fallback for older Ollama versions
ollama_models = data.get('items', [])
# Add Ollama models to the table with continuing numbers
# (after predefined models + LiteLLM models in numbering)
start_index = len(predefined_model_names) + len(litellm_model_names) + 1
for i, model in enumerate(ollama_models, start_index):
model_name = model.get('name', '')
model_size = model.get('size', 0)
# Convert size to human-readable format
size_str = ""
if model_size:
size_mb = model_size / (1024 * 1024)
if model_size < 1024 * 1024 * 1024:
size_str = f"{size_mb:.1f} MB"
else:
size_gb = size_mb / 1024
size_str = f"{size_gb:.1f} GB"
# Ollama models are free to use locally
model_description = "Local model"
if size_str:
model_description += f" ({size_str})"
model_table.add_row(
str(i),
model_name,
"Ollama",
"Local",
"Free",
"Free",
model_description
)
# Add to cached models for numeric selection
self.cached_models.append(model_name)
self.cached_model_numbers[str(i)] = model_name
except Exception: # pylint: disable=broad-except
# Ollama models (display from already loaded data)
if ollama_models_data:
start_index = len(predefined_model_names) + len(litellm_model_names) + 1
for i, model in enumerate(ollama_models_data, start_index):
model_name = model.get('name', '')
model_size = model.get('size', 0)
size_str = ""
if model_size:
size_mb = model_size / (1024 * 1024)
if model_size < 1024 * 1024 * 1024:
size_str = f"{size_mb:.1f} MB"
else:
size_gb = size_mb / 1024
size_str = f"{size_gb:.1f} GB"
model_description = "Local model"
if size_str:
model_description += f" ({size_str})"
model_table.add_row(
str(i),
model_name,
"Ollama",
"Local",
"Free",
"Free",
model_description
)
else: # pylint: disable=broad-except
# Add a note about Ollama if we couldn't fetch models
start_index = len(predefined_model_names) + len(litellm_model_names) + 1
model_table.add_row(
@ -436,6 +484,15 @@ class ModelShowCommand(Command):
if args: # If there are still args left, use as search term
search_term = args[0].lower()
# Load all models and update global cache for consistent numbering with /model
global _GLOBAL_MODEL_CACHE, _GLOBAL_MODEL_NUMBERS
all_model_names, ollama_models_data = load_all_available_models()
_GLOBAL_MODEL_CACHE = all_model_names
_GLOBAL_MODEL_NUMBERS = {
str(i): model_name
for i, model_name in enumerate(_GLOBAL_MODEL_CACHE, 1)
}
# Fetch model pricing data from LiteLLM GitHub repository
try:
with console.status(
@ -482,10 +539,53 @@ class ModelShowCommand(Command):
# Count models for summary
total_models = 0
displayed_models = 0
model_index = 1
# Process and display models
# 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:
model_index = _GLOBAL_MODEL_CACHE.index(model_name) + 1
except ValueError:
continue # Model not in cache, skip
total_models += 1
# Skip if showing only supported models and no function calling
@ -569,71 +669,46 @@ class ModelShowCommand(Command):
features_str
)
model_index += 1
# Now add Ollama models if available
try:
# Get Ollama models with a short timeout
api_base = get_ollama_api_base()
api_tags = f"{api_base.replace('/v1', '')}/api/tags"
ollama_response = requests.get(api_tags, timeout=1)
if ollama_response.status_code == 200:
ollama_data = ollama_response.json()
ollama_models = []
if 'models' in ollama_data:
ollama_models = ollama_data['models']
# Add Ollama models to the table (already loaded in global cache)
for model in ollama_models_data:
model_name = model.get('name', '')
# Skip if search term provided and not in model name
if search_term and search_term not in model_name.lower():
continue
# Find index from global cache
try:
model_index = _GLOBAL_MODEL_CACHE.index(model_name) + 1
except ValueError:
continue
total_models += 1
displayed_models += 1
model_size = model.get('size', 0)
size_str = ""
if model_size:
size_mb = model_size / (1024 * 1024)
if model_size < 1024 * 1024 * 1024:
size_str = f"{size_mb:.1f} MB"
else:
# Fallback for older Ollama versions
ollama_models = ollama_data.get('items', [])
# Add Ollama models to the table
for model in ollama_models:
model_name = model.get('name', '')
# Skip if search term provided and not in model name
if (search_term and
search_term not in model_name.lower()):
continue
total_models += 1
displayed_models += 1
model_size = model.get('size', 0)
# Convert size to human-readable format
size_str = ""
if model_size:
size_mb = model_size / (1024 * 1024)
if model_size < 1024 * 1024 * 1024:
size_str = f"{size_mb:.1f} MB"
else:
size_gb = size_mb / 1024
size_str = f"{size_gb:.1f} GB"
# Add row to table
model_description = "Local model"
if size_str:
model_description += f" ({size_str})"
model_table.add_row(
str(model_index),
model_name,
"Ollama",
"Varies",
"Free",
"Free",
model_description
)
model_index += 1
except Exception: # pylint: disable=broad-except
# Silently fail if Ollama is not available
# This is acceptable as Ollama is optional and we don't want to
# disrupt the user experience if it's not running
console.print(
"[dim]Ollama models not available[/dim]",
style="dim")
size_gb = size_mb / 1024
size_str = f"{size_gb:.1f} GB"
model_description = "Local model"
if size_str:
model_description += f" ({size_str})"
model_table.add_row(
str(model_index),
model_name,
"Ollama",
"Varies",
"Free",
"Free",
model_description
)
# Display the table
console.print(model_table)

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

2
src/cai/sdk/agents/function_schema.py
View File

@ -134,7 +134,7 @@ def _detect_docstring_style(doc: str) -> DocstringStyle:
@contextlib.contextmanager
def _suppress_griffe_logging():
# Supresses warnings about missing annotations for params
# Suppresses warnings about missing annotations for params
logger = logging.getLogger("griffe")
previous_level = logger.getEffectiveLevel()
logger.setLevel(logging.ERROR)

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:

26
src/cai/tools/command_and_control/sshpass.py
View File

@ -15,6 +15,7 @@ something that hasn't been seen in other cybersecurity frameworks yet (Feb 2025)
from cai.tools.common import run_command # pylint: disable=E0401 # noqa: E501
from cai.sdk.agents import function_tool
import shlex
@function_tool
def run_ssh_command_with_credentials(
@ -36,14 +37,25 @@ def run_ssh_command_with_credentials(
Returns:
str: Output from the remote command execution
"""
# Escape special characters in password and command to prevent shell injection
escaped_password = password.replace("'", "'\\''")
escaped_command = command.replace("'", "'\\''")
try:
port = int(port)
if port <= 0 or port > 65535:
return "port is not a valid integer"
except Exception:
return "port is not a valid integer"
# Escape special characters to prevent shell injection
quoted_password = shlex.quote(password)
quoted_username = shlex.quote(username)
quoted_host = shlex.quote(host)
quoted_command = shlex.quote(command)
port = str(port)
ssh_command = (
f"sshpass -p '{escaped_password}' "
f"sshpass -p {quoted_password} "
f"ssh -o StrictHostKeyChecking=no "
f"{username}@{host} -p {port} "
f"'{escaped_command}'"
f"{quoted_username}@{quoted_host} -p {port} "
f"{quoted_command}"
)
return run_command(ssh_command)

172
src/cai/tools/common.py
View File

@ -12,6 +12,7 @@ import time
import uuid
import sys
import shlex
import select
from wasabi import color # pylint: disable=import-error
from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer, cli_print_tool_output
@ -265,16 +266,22 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
self.is_running = False
return str(e)
def _read_output(self):
"""Read output from the process"""
"""Read output with non-blocking select"""
try:
while self.is_running and self.master is not None:
try:
# Check if process has exited before reading
if self.process and self.process.poll() is not None:
self.is_running = False
break
# Read raw output chunk from PTY (don't require newlines)
# Non-blocking check for data
ready, _, _ = select.select([self.master], [], [], 0.5)
if not ready:
if self.process and self.process.poll() is not None:
self.is_running = False
break
continue
output = os.read(self.master, 4096).decode('utf-8', errors='replace')
if output is not None and output != "":
@ -694,27 +701,43 @@ async def _run_local_async(command, stdout=False, timeout=100, stream=False, cal
# Don't add refresh_rate to tool_args as it affects command deduplication
# The refresh behavior is already handled by the streaming update logic
# Stream stdout in real-time
async for line in process.stdout:
line_str = line.decode('utf-8', errors='replace')
# Add to output collection
output_buffer.append(line_str)
buffer_size += 1
# Only update periodically to reduce UI refreshes
if buffer_size >= update_interval:
current_output = ''.join(output_buffer)
update_tool_streaming(tool_name, tool_args, current_output, call_id, token_info)
buffer_size = 0
# Stream stdout with idle detection
last_output = time.time()
while True:
if process.returncode is not None:
break
try:
line = await asyncio.wait_for(process.stdout.readline(), timeout=0.5)
if line:
output_buffer.append(line.decode('utf-8', errors='replace'))
buffer_size += 1
last_output = time.time()
if buffer_size >= update_interval:
update_tool_streaming(tool_name, tool_args, ''.join(output_buffer), call_id, token_info)
buffer_size = 0
else:
break
except asyncio.TimeoutError:
if time.time() - last_output > 10:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=1.0)
except asyncio.TimeoutError:
process.kill()
await process.wait()
output_buffer.append("\n[Terminated: idle 10s, likely waiting for input]")
break
# Wait for process to complete with timeout
try:
return_code = await asyncio.wait_for(process.wait(), timeout=timeout)
except asyncio.TimeoutError:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
# Wait for process to complete
if process.returncode is None:
try:
return_code = await asyncio.wait_for(process.wait(), timeout=timeout)
except asyncio.TimeoutError:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
else:
return_code = process.returncode
process_execution_time = time.time() - process_start_time
@ -743,7 +766,7 @@ async def _run_local_async(command, stdout=False, timeout=100, stream=False, cal
return final_output
else:
# Standard non-streaming async execution
# Non-streaming with idle detection
process = await asyncio.create_subprocess_shell(
command,
stdout=asyncio.subprocess.PIPE,
@ -751,15 +774,45 @@ async def _run_local_async(command, stdout=False, timeout=100, stream=False, cal
cwd=target_dir
)
try:
stdout_data, stderr_data = await asyncio.wait_for(
process.communicate(),
timeout=timeout
)
except asyncio.TimeoutError:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
stdout_chunks, stderr_chunks = [], []
last_output = time.time()
start = time.time()
while True:
if time.time() - start > timeout:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
if process.returncode is not None:
break
try:
out_task = asyncio.create_task(process.stdout.read(4096))
err_task = asyncio.create_task(process.stderr.read(4096))
done, pending = await asyncio.wait([out_task, err_task], timeout=0.5, return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
for task in done:
data = await task
if data:
(stdout_chunks if task == out_task else stderr_chunks).append(data)
last_output = time.time()
except asyncio.TimeoutError:
pass
if time.time() - last_output > 10:
try:
await asyncio.wait_for(process.wait(), timeout=0.1)
break
except asyncio.TimeoutError:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=1.0)
except asyncio.TimeoutError:
process.kill()
await process.wait()
stderr_chunks.append(b"\n[Terminated: idle 10s]")
break
stdout_data, stderr_data = b''.join(stdout_chunks), b''.join(stderr_chunks)
# Decode output
output = stdout_data.decode('utf-8', errors='replace') if stdout_data else ""
@ -968,26 +1021,43 @@ async def _run_docker_async(command, container_id, stdout=False, timeout=100, st
start_time = time.time()
# Read stdout line by line
async for line in process.stdout:
line_str = line.decode('utf-8', errors='replace')
output_buffer.append(line_str)
buffer_size += 1
# Only update periodically to reduce UI refreshes
if buffer_size >= update_interval:
# Show actual output as it's being collected
current_output = ''.join(output_buffer)
update_tool_streaming(tool_name, tool_args, current_output, call_id, token_info)
buffer_size = 0
# Read stdout with idle detection
last_output = time.time()
while True:
if process.returncode is not None:
break
try:
line = await asyncio.wait_for(process.stdout.readline(), timeout=0.5)
if line:
output_buffer.append(line.decode('utf-8', errors='replace'))
buffer_size += 1
last_output = time.time()
if buffer_size >= update_interval:
update_tool_streaming(tool_name, tool_args, ''.join(output_buffer), call_id, token_info)
buffer_size = 0
else:
break
except asyncio.TimeoutError:
if time.time() - last_output > 10:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=1.0)
except asyncio.TimeoutError:
process.kill()
await process.wait()
output_buffer.append("\n[Terminated: idle 10s]")
break
# Wait for process completion
try:
return_code = await asyncio.wait_for(process.wait(), timeout=timeout)
except asyncio.TimeoutError:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
if process.returncode is None:
try:
return_code = await asyncio.wait_for(process.wait(), timeout=timeout)
except asyncio.TimeoutError:
process.kill()
await process.wait()
raise subprocess.TimeoutExpired(command, timeout)
else:
return_code = process.returncode
execution_time = time.time() - start_time

2
src/cai/tools/web/headers.py
View File

@ -8,8 +8,10 @@ analysis, parameter inspection, and security vulnerability detection.
from urllib.parse import urlparse
import requests # pylint: disable=E0401
from cai.sdk.agents import function_tool
@function_tool(strict_mode=False)
def web_request_framework( # noqa: E501 # pylint: disable=too-many-arguments,too-many-locals,too-many-branches
url: str = "",
method: str = "GET",

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