mirror of https://github.com/aliasrobotics/cai.git
Default to alias0 and improve overall init. message
Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
parent
227afd6d12
commit
d7ee0c253c
|
|
@ -64,7 +64,7 @@ from dotenv import load_dotenv # pylint: disable=import-error # noqa: E501
|
|||
__path__ = pkgutil.extend_path(__path__, __name__)
|
||||
|
||||
# Get model from environment or use default
|
||||
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
model = os.getenv('CAI_MODEL', "alias0")
|
||||
|
||||
|
||||
PATTERNS = [
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ blueteam_agent = Agent(
|
|||
description="""Agent that specializes in system defense and security monitoring.
|
||||
Expert in cybersecurity protection and incident response.""",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
),
|
||||
tools=tools,
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ bug_bounter_agent = Agent(
|
|||
Expert in web security, API testing, and responsible disclosure.""",
|
||||
tools=tools,
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -184,7 +184,7 @@ class CodeAgent(Agent):
|
|||
def __init__( # pylint: disable=too-many-arguments,too-many-locals # noqa: E501
|
||||
self,
|
||||
name: str = "CodeAgent",
|
||||
model: str = "qwen2.5:14b",
|
||||
model: str = "alias0",
|
||||
instructions: Union[str, Callable[[], str]] = None,
|
||||
tools: List[Callable] = None,
|
||||
additional_authorized_imports: Optional[List[str]] = None,
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ dfir_agent = Agent(
|
|||
description="""Agent that specializes in Digital Forensics and Incident Response.
|
||||
Expert in investigation and analysis of digital evidence.""",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
),
|
||||
tools=tools,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, handoff
|
|||
from openai import AsyncOpenAI
|
||||
from cai.agents.one_tool import one_tool_agent
|
||||
|
||||
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
model = os.getenv('CAI_MODEL', "alias0")
|
||||
|
||||
# Create OpenAI client with fallback API key to prevent initialization errors
|
||||
# The actual API key should be set in environment variables or .env file
|
||||
|
|
@ -22,7 +22,7 @@ flag_discriminator = Agent(
|
|||
4. If you do not find a flag, call `ctf_agent` to continue investigating.
|
||||
""",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model="qwen2.5:14b" if os.getenv('CAI_MODEL') == "o3-mini" else model,
|
||||
model="alias0" if os.getenv('CAI_MODEL') == "o3-mini" else model,
|
||||
openai_client=AsyncOpenAI(api_key=api_key),
|
||||
),
|
||||
handoffs=[
|
||||
|
|
|
|||
|
|
@ -115,7 +115,7 @@ dns_smtp_agent = Agent(
|
|||
),
|
||||
tools=[check_mail_spoofing_vulnerability, execute_cli_command],
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -87,7 +87,7 @@ from cai.tools.misc.rag import add_to_memory_semantic, add_to_memory_episodic
|
|||
from cai.rag.vector_db import get_previous_memory
|
||||
|
||||
# Get model from environment or use default
|
||||
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
model = os.getenv('CAI_MODEL', "alias0")
|
||||
|
||||
|
||||
def get_previous_steps(query: str) -> str:
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ network_security_analyzer_agent = Agent(
|
|||
description="""Agent that specializes in network security analysis.
|
||||
Expert in monitoring, capturing, and analyzing network communications for security threats.""",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
),
|
||||
tools=tools,
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from cai.tools.reconnaissance.generic_linux_command import generic_linux_command
|
|||
from openai import AsyncOpenAI
|
||||
|
||||
# Get model from environment or use default
|
||||
model_name = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
model_name = os.getenv('CAI_MODEL', "alias0")
|
||||
|
||||
# NOTE: This is needed when using LiteLLM Proxy Server
|
||||
#
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error
|
|||
from cai.util import load_prompt_template
|
||||
|
||||
load_dotenv()
|
||||
model_name = os.getenv("CAI_MODEL", "qwen2.5:14b")
|
||||
model_name = os.getenv("CAI_MODEL", "alias0")
|
||||
# Prompts
|
||||
redteam_agent_system_prompt = load_prompt_template("prompts/system_red_team_agent.md")
|
||||
# Define tools list based on available API keys
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ replay_attack_agent = Agent(
|
|||
description="""Agent that specializes in network replay attacks and counteroffensive techniques.
|
||||
Expert in packet manipulation, traffic replay, and protocol exploitation.""",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
),
|
||||
tools=tools,
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ retester_agent = Agent(
|
|||
eliminating false positives.""",
|
||||
tools=tools,
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ thought_agent_system_prompt = load_prompt_template("prompts/system_thought_route
|
|||
thought_agent = Agent(
|
||||
name="ThoughtAgent",
|
||||
model=OpenAIChatCompletionsModel(
|
||||
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
|
||||
model=os.getenv('CAI_MODEL', "alias0"),
|
||||
openai_client=AsyncOpenAI(),
|
||||
),
|
||||
description="""Agent focused on analyzing and planning the next steps
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ Environment Variables
|
|||
within container (default: "true")
|
||||
|
||||
CAI_MODEL: Model to use for agents
|
||||
(default: "qwen2.5:14b")
|
||||
(default: "alias0")
|
||||
CAI_DEBUG: Set debug output level (default: "1")
|
||||
- 0: Only tool outputs
|
||||
- 1: Verbose debug output
|
||||
|
|
@ -182,7 +182,7 @@ set_tracing_disabled(True)
|
|||
|
||||
# llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini')
|
||||
# # llm_model=os.getenv('LLM_MODEL', 'claude-3-7')
|
||||
llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b')
|
||||
llm_model=os.getenv('LLM_MODEL', 'alias0')
|
||||
|
||||
|
||||
# For Qwen models, we need to skip system instructions as they're not supported
|
||||
|
|
@ -271,7 +271,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
turn_count = 0
|
||||
idle_time = 0
|
||||
console = Console()
|
||||
last_model = os.getenv('CAI_MODEL', 'qwen2.5:14b')
|
||||
last_model = os.getenv('CAI_MODEL', 'alias0')
|
||||
last_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent')
|
||||
parallel_count = int(os.getenv('CAI_PARALLEL', '1'))
|
||||
|
||||
|
|
@ -351,7 +351,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
idle_start_time = time.time()
|
||||
|
||||
# Check if model has changed and update if needed
|
||||
current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b')
|
||||
current_model = os.getenv('CAI_MODEL', 'alias0')
|
||||
if current_model != last_model and hasattr(agent, 'model'):
|
||||
# Update the model recursively for the agent and all handoff agents
|
||||
update_agent_models_recursively(agent, current_model)
|
||||
|
|
@ -926,7 +926,7 @@ def main():
|
|||
agent.model.suppress_final_output = False # Changed to False to show all agent messages
|
||||
|
||||
# Ensure the agent and all its handoff agents use the current model
|
||||
current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b')
|
||||
current_model = os.getenv('CAI_MODEL', 'alias0')
|
||||
update_agent_models_recursively(agent, current_model)
|
||||
|
||||
# Run the CLI with the selected agent
|
||||
|
|
|
|||
|
|
@ -46,7 +46,7 @@ ENV_VARS = {
|
|||
6: {
|
||||
"name": "CAI_MODEL",
|
||||
"description": "Model to use for agents",
|
||||
"default": "qwen2.5:14b"
|
||||
"default": "alias0"
|
||||
},
|
||||
7: {
|
||||
"name": "CAI_DEBUG",
|
||||
|
|
|
|||
|
|
@ -51,7 +51,7 @@ class FlushCommand(Command):
|
|||
|
||||
if client and hasattr(client, 'interaction_input_tokens') and hasattr(
|
||||
client, 'total_input_tokens'):
|
||||
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
model = os.getenv('CAI_MODEL', "alias0")
|
||||
input_tokens = client.interaction_input_tokens if hasattr(
|
||||
client, 'interaction_input_tokens') else 0
|
||||
total_tokens = client.total_input_tokens if hasattr(
|
||||
|
|
|
|||
|
|
@ -290,7 +290,7 @@ def display_quick_guide(console: Console):
|
|||
)
|
||||
|
||||
# Get current environment variable values
|
||||
current_model = os.getenv('CAI_MODEL', "qwen2.5:14b")
|
||||
current_model = os.getenv('CAI_MODEL', "alias0")
|
||||
current_agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent")
|
||||
|
||||
config_text = Text.assemble(
|
||||
|
|
@ -330,16 +330,23 @@ def display_quick_guide(console: Console):
|
|||
)
|
||||
|
||||
context_tip = Panel(
|
||||
"As security exercises progress, LLM quality may\n"
|
||||
"degrade, especially if progress stalls.\n\n"
|
||||
"It's often better to clear the context window\n"
|
||||
"or restart CAI rather than waiting until\n"
|
||||
"context usage reaches 100%.\n\n"
|
||||
"When context exceeds 80%, follow these steps:\n"
|
||||
"1. CAI> Dump your memory and findings in current scenario in findings.txt\n"
|
||||
"2. CAI> /flush\n"
|
||||
"3. CAI> Analyze findings.txt, and continue exercise with target: ...",
|
||||
title="[bold yellow]Performance Tip[/bold yellow]",
|
||||
Text.assemble(
|
||||
"For optimal cybersecurity AI performance, use\n",
|
||||
("alias0", "bold green"),
|
||||
" - specifically designed for cybersecurity\n"
|
||||
"tasks with superior domain knowledge.\n\n",
|
||||
("alias0", "bold green"),
|
||||
" outperforms general-purpose models in:\n",
|
||||
"• Vulnerability assessment\n",
|
||||
"• Penetration testing and bug bounty\n",
|
||||
"• Security analysis\n",
|
||||
"• Threat detection\n\n",
|
||||
"Learn more about ",
|
||||
("alias0", "bold green"),
|
||||
" and its privacy-first approach:\n",
|
||||
("https://news.aliasrobotics.com/alias0-a-privacy-first-cybersecurity-ai/", "blue underline")
|
||||
),
|
||||
title="[bold yellow]Cybersecurity Model Tip[/bold yellow]",
|
||||
border_style="yellow",
|
||||
padding=(1, 2),
|
||||
title_align="center"
|
||||
|
|
|
|||
|
|
@ -2236,7 +2236,7 @@ class OpenAIChatCompletionsModel(Model):
|
|||
elif "gemini" in model_str:
|
||||
kwargs.pop("parallel_tool_calls", None)
|
||||
elif "qwen" in model_str or ":" in model_str:
|
||||
# Handle Ollama-served models with custom formats (e.g., qwen2.5:14b)
|
||||
# Handle Ollama-served models with custom formats (e.g., alias0)
|
||||
# These typically need the Ollama provider
|
||||
litellm.drop_params = True
|
||||
kwargs.pop("parallel_tool_calls", None)
|
||||
|
|
|
|||
|
|
@ -1805,7 +1805,7 @@ def update_agent_streaming_content(context, text_delta, token_stats=None):
|
|||
footer_stats.append(f"${session_total_cost:.4f}", style="bold magenta")
|
||||
|
||||
# Add context usage indicator
|
||||
model_name = context.get("model", os.environ.get('CAI_MODEL', 'qwen2.5:14b'))
|
||||
model_name = context.get("model", os.environ.get('CAI_MODEL', 'alias0'))
|
||||
context_pct = input_tokens / get_model_input_tokens(model_name) * 100
|
||||
if context_pct < 50:
|
||||
indicator = "🟩"
|
||||
|
|
@ -3109,7 +3109,7 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None,
|
|||
|
||||
# Calculate cost if not provided
|
||||
if not interaction_cost and input_tokens > 0:
|
||||
model_name = token_info.get('model', os.environ.get('CAI_MODEL', 'qwen2.5:14b'))
|
||||
model_name = token_info.get('model', os.environ.get('CAI_MODEL', 'alias0'))
|
||||
interaction_cost = calculate_model_cost(model_name, input_tokens, output_tokens)
|
||||
|
||||
# Add compact token info to output
|
||||
|
|
|
|||
Loading…
Reference in New Issue