From be455aba80ba4a1502e70bde2caf7f93ef4188e8 Mon Sep 17 00:00:00 2001 From: "Cristian \"SkyH34D\" Franco" <58003229+SkyH34D@users.noreply.github.com> Date: Sun, 7 Dec 2025 12:06:23 +0100 Subject: [PATCH 1/8] Fix typo in README.md for 'Agents' (#366) Correcting typo and adjusting diagram --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 5e8f706a..21658c7b 100644 --- a/README.md +++ b/README.md @@ -545,21 +545,21 @@ OLLAMA_API_BASE="https://custom-openai-proxy.com/v1" cai ## :triangular_ruler: Architecture: -CAI focuses on making cybersecurity agent **coordination** and **execution** lightweight, highly controllable, and useful for humans. To do so it builds upon 8 pillars: `Agent`s, `Tools`, `Handoffs`, `Patterns`, `Turns`, `Tracing`, `Guardrails` and `HITL`. +CAI focuses on making cybersecurity agent **coordination** and **execution** lightweight, highly controllable, and useful for humans. To do so it builds upon 8 pillars: `Agents`, `Tools`, `Handoffs`, `Patterns`, `Turns`, `Tracing`, `Guardrails` and `HITL`. ``` ┌───────────────┐ ┌───────────┐ - │ HITL │◀─────────▶│ Turns │ + │ HITL │◀────────▶│ Turns │ └───────┬───────┘ └───────────┘ │ ▼ ┌───────────┐ ┌───────────┐ ┌───────────┐ ┌───────────┐ -│ Patterns │◀─────▶│ Handoffs │◀────▶ │ Agents │◀────▶│ LLMs │ +│ Patterns │◀────▶│ Handoffs │◀───▶ │ Agents │◀───▶│ LLMs │ └───────────┘ └─────┬─────┘ └─────┬─────┘ └───────────┘ │ │ │ ▼ ┌────────────┐ ┌────┴──────┐ ┌───────────┐ ┌────────────┐ -│ Extensions │◀─────▶│ Tracing │ │ Tools │◀───▶│ Guardrails │ +│ Extensions │◀────▶│ Tracing │ │ Tools │◀──▶│ Guardrails │ └────────────┘ └───────────┘ └───────────┘ └────────────┘ │ ┌─────────────┬─────┴────┬─────────────┐ From 2ed9951cf2fdd5e2e3830c7a5b08fa29b467a435 Mon Sep 17 00:00:00 2001 From: Didier Durand <2927957+didier-durand@users.noreply.github.com> Date: Mon, 8 Dec 2025 09:16:53 +0100 Subject: [PATCH 2/8] [Doc]: fixing typos in different files (#364) --- docs/cai_architecture.md | 16 ++++++++-------- docs/cai_pro_contact.md | 2 +- docs/cli/commands_reference.md | 2 +- docs/context.md | 2 +- src/cai/agents/__init__.py | 2 +- src/cai/sdk/agents/function_schema.py | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) 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/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/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) From d37584f4571bb97b904aa51a0c925ac89114c05a Mon Sep 17 00:00:00 2001 From: Cristi Date: Mon, 8 Dec 2025 12:05:51 +0200 Subject: [PATCH 3/8] Add web_request_framework tool to web_pentester (#358) Co-authored-by: root --- src/cai/agents/web_pentester.py | 2 ++ src/cai/tools/web/headers.py | 2 ++ 2 files changed, 4 insertions(+) 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/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", From 09ccb6e0baccf56c40e6cb429c698750843a999c Mon Sep 17 00:00:00 2001 From: Edoardo Ottavianelli Date: Tue, 9 Dec 2025 16:36:22 +0100 Subject: [PATCH 4/8] Merge commit from fork --- src/cai/tools/command_and_control/sshpass.py | 26 ++++++++++++++------ 1 file changed, 19 insertions(+), 7 deletions(-) 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) From e932e3085013e832d360d5fa965f65af209b31b5 Mon Sep 17 00:00:00 2001 From: pzabalegui Date: Tue, 9 Dec 2025 18:42:12 +0100 Subject: [PATCH 5/8] fix: prevent indefinite blocking on interactive commands (#370) Add idle detection to prevent CAI from freezing when executing interactive tools (cat, python -i, gdb, radare2) that wait for user input indefinitely. - Add non-blocking I/O with asyncio.wait_for() timeouts - Add 10s idle timeout to terminate idle processes - Apply to streaming and non-streaming execution paths - Add select.select() to ShellSession._read_output() Resolves blocking I/O issue where interactive tools would hang CAI for the full timeout period (100s). --- src/cai/tools/common.py | 172 ++++++++++++++++++++++++++++------------ 1 file changed, 121 insertions(+), 51 deletions(-) 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 From 7c2267d2b9bd2dbaec4b02d03ecd1781b51125a7 Mon Sep 17 00:00:00 2001 From: pzabalegui Date: Wed, 10 Dec 2025 12:41:35 +0100 Subject: [PATCH 6/8] =?UTF-8?q?fix:=20consistent=20model=20numbering=20bet?= =?UTF-8?q?ween=20/model=20and=20/model-show=20for=20Ol=E2=80=A6=20(#371)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix: consistent model numbering between /model and /model-show for Ollama When using OLLAMA_API_BASE, model selection by number was incorrect. /model-show displayed Ollama models with certain numbers, but /model selected different models from LiteLLM instead. This happened because both commands used separate caches with inconsistent numbering. Now they share a global cache that loads models in consistent order: predefined → LiteLLM → Ollama. Fixes model number mismatch when selecting Ollama models by number. --- src/cai/repl/commands/model.py | 279 ++++++++++++++++----------------- 1 file changed, 137 insertions(+), 142 deletions(-) diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index a0a64994..a16bc5d4 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -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. @@ -155,6 +159,43 @@ 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: + litellm_names = sorted(response.json().keys()) + except Exception: # pylint: disable=broad-except + pass + + # Ollama models + ollama_data = [] + ollama_names = [] + try: + api_base = get_ollama_api_base() + response = requests.get(f"{api_base.replace('/v1', '')}/api/tags", timeout=1) + 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 +237,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 +305,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 +443,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 +498,14 @@ class ModelShowCommand(Command): # Count models for summary total_models = 0 displayed_models = 0 - model_index = 1 - # Process and display models + # Process and display 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 +589,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) From da1c6202e5857494a1968937c60c5030467dcabd Mon Sep 17 00:00:00 2001 From: pzabalegui Date: Wed, 10 Dec 2025 15:42:48 +0100 Subject: [PATCH 7/8] 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. --- docs/providers/ollama.md | 27 ++++++ docs/providers/ollama_cloud.md | 79 ++++++++++++++++ src/cai/repl/commands/completer.py | 12 ++- src/cai/repl/commands/model.py | 90 +++++++++++++++++-- src/cai/repl/ui/banner.py | 15 +++- src/cai/repl/ui/toolbar.py | 17 +++- .../agents/models/openai_chatcompletions.py | 63 ++++++++++++- src/cai/util.py | 44 ++++++++- 8 files changed, 329 insertions(+), 18 deletions(-) create mode 100644 docs/providers/ollama_cloud.md 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/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 a16bc5d4..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() @@ -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: diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py index 0050b992..a5e40140 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/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/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") From d086cc4871ce7ce09496a8860fbafd46e0dc28fc Mon Sep 17 00:00:00 2001 From: "@uayucar" Date: Fri, 19 Dec 2025 12:39:14 +0100 Subject: [PATCH 8/8] new version with Ollama Cloud --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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"