diff --git a/docs/cai_architecture.md b/docs/cai_architecture.md index 64f2daeb..b62e0b41 100644 --- a/docs/cai_architecture.md +++ b/docs/cai_architecture.md @@ -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) diff --git a/docs/cai_pro_contact.md b/docs/cai_pro_contact.md index 141fb2ac..a492dc6b 100644 --- a/docs/cai_pro_contact.md +++ b/docs/cai_pro_contact.md @@ -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) diff --git a/docs/cli/commands_reference.md b/docs/cli/commands_reference.md index 4cc1e12d..ea8ec9cf 100644 --- a/docs/cli/commands_reference.md +++ b/docs/cli/commands_reference.md @@ -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. diff --git a/docs/context.md b/docs/context.md index 5454c0ad..de5e8e22 100644 --- a/docs/context.md +++ b/docs/context.md @@ -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. diff --git a/docs/providers/ollama.md b/docs/providers/ollama.md index 1ff39758..34f32aec 100644 --- a/docs/providers/ollama.md +++ b/docs/providers/ollama.md @@ -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. diff --git a/docs/providers/ollama_cloud.md b/docs/providers/ollama_cloud.md new file mode 100644 index 00000000..6498731e --- /dev/null +++ b/docs/providers/ollama_cloud.md @@ -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) diff --git a/pyproject.toml b/pyproject.toml index d72bb1df..c09a0784 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/cai/agents/__init__.py b/src/cai/agents/__init__.py index acd914f4..932d9147 100644 --- a/src/cai/agents/__init__.py +++ b/src/cai/agents/__init__.py @@ -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.* | diff --git a/src/cai/agents/web_pentester.py b/src/cai/agents/web_pentester.py index a78b08d3..66bbfbd5 100644 --- a/src/cai/agents/web_pentester.py +++ b/src/cai/agents/web_pentester.py @@ -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 diff --git a/src/cai/repl/commands/completer.py b/src/cai/repl/commands/completer.py index f655d9c7..0c313335 100644 --- a/src/cai/repl/commands/completer.py +++ b/src/cai/repl/commands/completer.py @@ -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() diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index a0a64994..3afdf6fc 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -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) diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py index de288a16..f509cb1b 100644 --- a/src/cai/repl/ui/banner.py +++ b/src/cai/repl/ui/banner.py @@ -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 ) diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index 06b1e70f..fd6fd282 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -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() diff --git a/src/cai/sdk/agents/function_schema.py b/src/cai/sdk/agents/function_schema.py index 08fc6008..89c0b938 100644 --- a/src/cai/sdk/agents/function_schema.py +++ b/src/cai/sdk/agents/function_schema.py @@ -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) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 7e88f438..8931edd6 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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: diff --git a/src/cai/tools/command_and_control/sshpass.py b/src/cai/tools/command_and_control/sshpass.py index 6dc3d217..251b88a9 100644 --- a/src/cai/tools/command_and_control/sshpass.py +++ b/src/cai/tools/command_and_control/sshpass.py @@ -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) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 97db395a..af76f411 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -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 diff --git a/src/cai/tools/web/headers.py b/src/cai/tools/web/headers.py index bcca232f..004e6e0f 100644 --- a/src/cai/tools/web/headers.py +++ b/src/cai/tools/web/headers.py @@ -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", diff --git a/src/cai/util.py b/src/cai/util.py index a935e944..b5acca68 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -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")