From 394b0f400cf367c659eb32533cec35f6515ce1d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 23 May 2025 14:48:43 +0000 Subject: [PATCH] Hardenize the dotenv loading MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- src/cai/agents/flag_discriminator.py | 6 +++++- src/cai/agents/one_tool.py | 6 +++++- src/cai/cli.py | 14 +++++++++----- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/cai/agents/flag_discriminator.py b/src/cai/agents/flag_discriminator.py index 05e5e0ab..f1228810 100644 --- a/src/cai/agents/flag_discriminator.py +++ b/src/cai/agents/flag_discriminator.py @@ -8,6 +8,10 @@ from cai.agents.one_tool import one_tool_agent model = os.getenv('CAI_MODEL', "qwen2.5:14b") +# Create OpenAI client with fallback API key to prevent initialization errors +# The actual API key should be set in environment variables or .env file +api_key = os.getenv('OPENAI_API_KEY', 'sk-placeholder-key-for-local-models') + flag_discriminator = Agent( name="Flag discriminator", description="Agent focused on extracting the flag from the output", @@ -19,7 +23,7 @@ flag_discriminator = Agent( """, model=OpenAIChatCompletionsModel( model="qwen2.5:14b" if os.getenv('CAI_MODEL') == "o3-mini" else model, - openai_client=AsyncOpenAI(), + openai_client=AsyncOpenAI(api_key=api_key), ), handoffs=[ handoff( diff --git a/src/cai/agents/one_tool.py b/src/cai/agents/one_tool.py index 6dd09aad..bc9f4c58 100644 --- a/src/cai/agents/one_tool.py +++ b/src/cai/agents/one_tool.py @@ -50,6 +50,10 @@ instructions = """You are a Cybersecurity expert Leader facing a CTF """ +# Create OpenAI client with fallback API key to prevent initialization errors +# The actual API key should be set in environment variables or .env file +api_key = os.getenv('OPENAI_API_KEY', 'sk-placeholder-key-for-local-models') + one_tool_agent = Agent( name="CTF agent", description="""Agent focused on conquering security challenges using generic linux commands @@ -60,7 +64,7 @@ one_tool_agent = Agent( ], model=OpenAIChatCompletionsModel( model=model_name, - openai_client=AsyncOpenAI(), + openai_client=AsyncOpenAI(api_key=api_key), ) ) diff --git a/src/cai/cli.py b/src/cai/cli.py index 469aa4c2..1258b77d 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -107,10 +107,13 @@ Usage Examples: CAI_MODEL="gpt-4o" CAI_PARALLEL="3" cai """ +# Load environment variables from .env file FIRST, before any imports import os +from dotenv import load_dotenv +load_dotenv() + import time import asyncio -from dotenv import load_dotenv from rich.console import Console from rich.panel import Panel @@ -164,9 +167,6 @@ if is_pentestperf_available() and os.getenv('CTF_NAME', None): container_id = "" os.environ['CAI_ACTIVE_CONTAINER'] = container_id -# Load environment variables from .env file -load_dotenv() - # NOTE: This is needed when using LiteLLM Proxy Server # # external_client = AsyncOpenAI( @@ -189,12 +189,16 @@ llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b') # For Qwen models, we need to skip system instructions as they're not supported instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant" +# Create OpenAI client with fallback API key to prevent initialization errors +# The actual API key should be set in environment variables or .env file +api_key = os.getenv('OPENAI_API_KEY', 'sk-placeholder-key-for-local-models') + agent = Agent( name="Assistant", instructions=instructions, model=OpenAIChatCompletionsModel( model=llm_model, - openai_client=AsyncOpenAI() # original OpenAI servers + openai_client=AsyncOpenAI(api_key=api_key) # original OpenAI servers # openai_client = external_client # LiteLLM Proxy Server ) )