From 244af4cf17a18ce23fc364777c075d86450a8220 Mon Sep 17 00:00:00 2001 From: toadlyBroodle Date: Sun, 4 Jan 2026 22:43:14 +0900 Subject: [PATCH] feat: add multi-model support and environment variable API keys Features: - Add separate planning_model and implementation_model for all providers - Support GOOGLE_API_KEY, ANTHROPIC_API_KEY, OPENAI_API_KEY env vars - Environment variables take precedence over mcp_agent.secrets.yaml - Add .env.example template for easy setup Refactoring: - Add shared get_api_keys() and load_api_config() in llm_utils.py - DRY up duplicated config loading code in workflow files - Use lazy imports for anthropic/openai to make them optional - Remove hardcoded brave search server, use config-based search Config changes: - Add planning_model/implementation_model options for google, anthropic, openai - Add mcp_agent.secrets.yaml to .gitignore (use .env instead) - Ignore deepcode_lab/, and logs/ --- .gitignore | 7 + mcp_agent.config.yaml | 17 +- mcp_agent.secrets.yaml | 7 + utils/llm_utils.py | 188 +++++++++++++----- workflows/agent_orchestration_engine.py | 16 +- workflows/code_implementation_workflow.py | 25 ++- .../code_implementation_workflow_index.py | 25 ++- 7 files changed, 210 insertions(+), 75 deletions(-) diff --git a/.gitignore b/.gitignore index 0f34269d..b9ed3728 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ site/ *.logfire *.coverage/ log/ +logs/ # Caches .cache/ @@ -64,3 +65,9 @@ run_indexer_with_filtering.py # Cline files memory-bank/ + +# project files +deepcode_lab/ + +# secrets (use .env or environment variables instead) +mcp_agent.secrets.yaml \ No newline at end of file diff --git a/mcp_agent.config.yaml b/mcp_agent.config.yaml index fd8a6e93..864ab3fa 100644 --- a/mcp_agent.config.yaml +++ b/mcp_agent.config.yaml @@ -1,6 +1,6 @@ $schema: ./schema/mcp-agent.config.schema.json anthropic: null -default_search_server: brave +default_search_server: filesystem document_segmentation: enabled: false size_threshold_chars: 50000 @@ -40,10 +40,12 @@ mcp: BRAVE_API_KEY: '' filesystem: # macos and linux should use this + # Note: "No valid root directories" warning is harmless - connection still works args: - -y - '@modelcontextprotocol/server-filesystem' - . + - ./deepcode_lab command: npx # windows should use this @@ -116,11 +118,22 @@ openai: max_tokens_policy: adaptive retry_max_tokens: 32768 -# Configuration for Google AI (Gemini) +# Provider configurations +# default_model is used by mcp_agent for planning/analysis phases +# implementation_model is used by code_implementation_workflow for code generation google: default_model: "gemini-3-pro-preview" + planning_model: "gemini-3-pro-preview" + implementation_model: "gemini-2.5-flash" anthropic: default_model: "claude-sonnet-4.5" + planning_model: "claude-sonnet-4.5" + implementation_model: "claude-sonnet-3.5" + +openai: + default_model: "o3-mini" + planning_model: "o3-mini" + implementation_model: "gpt-4o" planning_mode: traditional diff --git a/mcp_agent.secrets.yaml b/mcp_agent.secrets.yaml index 6c797f1f..79369cdd 100644 --- a/mcp_agent.secrets.yaml +++ b/mcp_agent.secrets.yaml @@ -1,3 +1,10 @@ +# API keys for LLM providers +# You can either fill these in directly, or use environment variables: +# - GOOGLE_API_KEY / GEMINI_API_KEY +# - ANTHROPIC_API_KEY +# - OPENAI_API_KEY +# Environment variables take precedence over values in this file. + openai: api_key: "" base_url: "" diff --git a/utils/llm_utils.py b/utils/llm_utils.py index f8bbccdd..23026d51 100644 --- a/utils/llm_utils.py +++ b/utils/llm_utils.py @@ -9,10 +9,89 @@ import yaml from typing import Any, Type, Dict, Tuple -# Import LLM classes -from mcp_agent.workflows.llm.augmented_llm_anthropic import AnthropicAugmentedLLM -from mcp_agent.workflows.llm.augmented_llm_openai import OpenAIAugmentedLLM -from mcp_agent.workflows.llm.augmented_llm_google import GoogleAugmentedLLM + +def get_api_keys(secrets_path: str = "mcp_agent.secrets.yaml") -> Dict[str, str]: + """ + Get API keys from environment variables or secrets file. + + Environment variables take precedence: + - GOOGLE_API_KEY or GEMINI_API_KEY + - ANTHROPIC_API_KEY + - OPENAI_API_KEY + + Args: + secrets_path: Path to the secrets YAML file + + Returns: + Dict with 'google', 'anthropic', 'openai' keys + """ + secrets = {} + if os.path.exists(secrets_path): + with open(secrets_path, "r", encoding="utf-8") as f: + secrets = yaml.safe_load(f) or {} + + return { + "google": ( + os.environ.get("GOOGLE_API_KEY") or + os.environ.get("GEMINI_API_KEY") or + secrets.get("google", {}).get("api_key", "") + ).strip(), + "anthropic": ( + os.environ.get("ANTHROPIC_API_KEY") or + secrets.get("anthropic", {}).get("api_key", "") + ).strip(), + "openai": ( + os.environ.get("OPENAI_API_KEY") or + secrets.get("openai", {}).get("api_key", "") + ).strip(), + } + + +def load_api_config(secrets_path: str = "mcp_agent.secrets.yaml") -> Dict[str, Any]: + """ + Load API configuration with environment variable override. + + Environment variables take precedence over YAML values: + - GOOGLE_API_KEY or GEMINI_API_KEY + - ANTHROPIC_API_KEY + - OPENAI_API_KEY + + Args: + secrets_path: Path to the secrets YAML file + + Returns: + Dict with provider configs including api_key values + """ + # Load base config from YAML + config = {} + if os.path.exists(secrets_path): + with open(secrets_path, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) or {} + + # Get keys with env var override + keys = get_api_keys(secrets_path) + + # Merge into config structure + for provider, key in keys.items(): + if key: + config.setdefault(provider, {})["api_key"] = key + + return config + + +def _get_llm_class(provider: str) -> Type[Any]: + """Lazily import and return the LLM class for a given provider.""" + if provider == "anthropic": + from mcp_agent.workflows.llm.augmented_llm_anthropic import AnthropicAugmentedLLM + return AnthropicAugmentedLLM + elif provider == "openai": + from mcp_agent.workflows.llm.augmented_llm_openai import OpenAIAugmentedLLM + return OpenAIAugmentedLLM + elif provider == "google": + from mcp_agent.workflows.llm.augmented_llm_google import GoogleAugmentedLLM + return GoogleAugmentedLLM + else: + raise ValueError(f"Unknown provider: {provider}") def get_preferred_llm_class(config_path: str = "mcp_agent.secrets.yaml") -> Type[Any]: @@ -31,18 +110,11 @@ def get_preferred_llm_class(config_path: str = "mcp_agent.secrets.yaml") -> Type class: The preferred LLM class """ try: - # Read API keys from secrets file - if not os.path.exists(config_path): - print(f"πŸ€– Config file {config_path} not found, using OpenAIAugmentedLLM") - return OpenAIAugmentedLLM - - with open(config_path, "r", encoding="utf-8") as f: - secrets = yaml.safe_load(f) - - # Get API keys - anthropic_key = secrets.get("anthropic", {}).get("api_key", "").strip() - google_key = secrets.get("google", {}).get("api_key", "").strip() - openai_key = secrets.get("openai", {}).get("api_key", "").strip() + # Get API keys with environment variable override + keys = get_api_keys(config_path) + google_key = keys["google"] + anthropic_key = keys["anthropic"] + openai_key = keys["openai"] # Read user preference from main config main_config_path = "mcp_agent.config.yaml" @@ -52,42 +124,38 @@ def get_preferred_llm_class(config_path: str = "mcp_agent.secrets.yaml") -> Type main_config = yaml.safe_load(f) preferred_provider = main_config.get("llm_provider", "").strip().lower() - # Map of providers to their classes and keys - provider_map = { - "anthropic": ( - AnthropicAugmentedLLM, - anthropic_key, - "AnthropicAugmentedLLM", - ), - "google": (GoogleAugmentedLLM, google_key, "GoogleAugmentedLLM"), - "openai": (OpenAIAugmentedLLM, openai_key, "OpenAIAugmentedLLM"), + # Map of providers to their keys and class names + provider_keys = { + "anthropic": (anthropic_key, "AnthropicAugmentedLLM"), + "google": (google_key, "GoogleAugmentedLLM"), + "openai": (openai_key, "OpenAIAugmentedLLM"), } # Try user's preferred provider first - if preferred_provider and preferred_provider in provider_map: - llm_class, api_key, class_name = provider_map[preferred_provider] + if preferred_provider and preferred_provider in provider_keys: + api_key, class_name = provider_keys[preferred_provider] if api_key: print(f"πŸ€– Using {class_name} (user preference: {preferred_provider})") - return llm_class + return _get_llm_class(preferred_provider) else: print( f"⚠️ Preferred provider '{preferred_provider}' has no API key, checking alternatives..." ) # Fallback: try providers in order of availability - for provider, (llm_class, api_key, class_name) in provider_map.items(): + for provider, (api_key, class_name) in provider_keys.items(): if api_key: print(f"πŸ€– Using {class_name} ({provider} API key found)") - return llm_class + return _get_llm_class(provider) - # No API keys found - print("⚠️ No API keys configured, falling back to OpenAIAugmentedLLM") - return OpenAIAugmentedLLM + # No API keys found - default to google + print("⚠️ No API keys configured, falling back to GoogleAugmentedLLM") + return _get_llm_class("google") except Exception as e: print(f"πŸ€– Error reading config file {config_path}: {e}") - print("πŸ€– Falling back to OpenAIAugmentedLLM") - return OpenAIAugmentedLLM + print("πŸ€– Falling back to GoogleAugmentedLLM") + return _get_llm_class("google") def get_token_limits(config_path: str = "mcp_agent.config.yaml") -> Tuple[int, int]: @@ -138,7 +206,8 @@ def get_default_models(config_path: str = "mcp_agent.config.yaml"): config_path: Path to the configuration file Returns: - dict: Dictionary with 'anthropic', 'openai', and 'google' default models + dict: Dictionary with 'anthropic', 'openai', 'google' default models, + plus 'google_planning' and 'google_implementation' for phase-specific models """ try: if os.path.exists(config_path): @@ -155,27 +224,54 @@ def get_default_models(config_path: str = "mcp_agent.config.yaml"): ) openai_model = openai_config.get("default_model", "o3-mini") google_model = google_config.get("default_model", "gemini-2.0-flash") + + # Phase-specific models (fall back to default if not specified) + # Google + google_planning = google_config.get("planning_model", google_model) + google_implementation = google_config.get("implementation_model", google_model) + # Anthropic + anthropic_planning = anthropic_config.get("planning_model", anthropic_model) + anthropic_implementation = anthropic_config.get("implementation_model", anthropic_model) + # OpenAI + openai_planning = openai_config.get("planning_model", openai_model) + openai_implementation = openai_config.get("implementation_model", openai_model) return { "anthropic": anthropic_model, "openai": openai_model, "google": google_model, + "google_planning": google_planning, + "google_implementation": google_implementation, + "anthropic_planning": anthropic_planning, + "anthropic_implementation": anthropic_implementation, + "openai_planning": openai_planning, + "openai_implementation": openai_implementation, } else: print(f"Config file {config_path} not found, using default models") - return { - "anthropic": "claude-sonnet-4-20250514", - "openai": "o3-mini", - "google": "gemini-2.0-flash", - } + return _get_fallback_models() except Exception as e: print(f"❌Error reading config file {config_path}: {e}") - return { - "anthropic": "claude-sonnet-4-20250514", - "openai": "o3-mini", - "google": "gemini-2.0-flash", - } + return _get_fallback_models() + + +def _get_fallback_models(): + """Return fallback model configuration when config file is unavailable.""" + google = "gemini-2.0-flash" + anthropic = "claude-sonnet-4-20250514" + openai = "o3-mini" + return { + "google": google, + "google_planning": google, + "google_implementation": google, + "anthropic": anthropic, + "anthropic_planning": anthropic, + "anthropic_implementation": anthropic, + "openai": openai, + "openai_planning": openai, + "openai_implementation": openai, + } def get_document_segmentation_config( diff --git a/workflows/agent_orchestration_engine.py b/workflows/agent_orchestration_engine.py index 68af299b..056f44b8 100644 --- a/workflows/agent_orchestration_engine.py +++ b/workflows/agent_orchestration_engine.py @@ -690,23 +690,17 @@ async def run_code_analyzer( prompts = get_adaptive_prompts(use_segmentation) if paper_content: + # When paper content is already loaded, agents don't need search tools agent_config = { "concept_analysis": [], - "algorithm_analysis": ["brave"], - "code_planner": [ - "brave" - ], # Empty list instead of None - code planner doesn't need tools when paper content is provided + "algorithm_analysis": search_server_names, + "code_planner": search_server_names, } - # agent_config = { - # "concept_analysis": [], - # "algorithm_analysis": [], - # "code_planner": [], # Empty list instead of None - code planner doesn't need tools when paper content is provided - # } else: agent_config = { "concept_analysis": ["filesystem"], - "algorithm_analysis": ["brave", "filesystem"], - "code_planner": ["brave", "filesystem"], + "algorithm_analysis": search_server_names + ["filesystem"], + "code_planner": search_server_names + ["filesystem"], } print(f" Agent configurations: {agent_config}") diff --git a/workflows/code_implementation_workflow.py b/workflows/code_implementation_workflow.py index c2a7b29f..77a5829a 100644 --- a/workflows/code_implementation_workflow.py +++ b/workflows/code_implementation_workflow.py @@ -33,7 +33,7 @@ from workflows.agents import CodeImplementationAgent from workflows.agents.memory_agent_concise import ConciseMemoryAgent from config.mcp_tool_definitions import get_mcp_tools -from utils.llm_utils import get_preferred_llm_class, get_default_models +from utils.llm_utils import get_preferred_llm_class, get_default_models, load_api_config # DialogueLogger removed - no longer needed @@ -61,10 +61,9 @@ def __init__(self, config_path: str = "mcp_agent.secrets.yaml"): ) def _load_api_config(self) -> Dict[str, Any]: - """Load API configuration from YAML file""" + """Load API configuration with environment variable override.""" try: - with open(self.config_path, "r", encoding="utf-8") as f: - return yaml.safe_load(f) + return load_api_config(self.config_path) except Exception as e: raise Exception(f"Failed to load API config: {e}") @@ -682,8 +681,11 @@ async def _call_anthropic_with_tools( ] try: + # Use implementation-specific model for code generation + impl_model = self.default_models.get("anthropic_implementation", self.default_models["anthropic"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") response = await client.messages.create( - model=self.default_models["anthropic"], + model=impl_model, system=system_message, messages=validated_messages, tools=tools, @@ -784,8 +786,11 @@ async def _call_google_with_tools( try: # Google Gemini API call using the native SDK # client is google.genai.Client instance + # Use implementation-specific model for code generation + impl_model = self.default_models.get("google_implementation", self.default_models["google"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") response = await client.aio.models.generate_content( - model=self.default_models["google"], + model=impl_model, contents=gemini_messages, config=config, ) @@ -1008,12 +1013,16 @@ async def _call_openai_with_tools( max_retries = 3 retry_delay = 2 # seconds + # Use implementation-specific model for code generation + impl_model = self.default_models.get("openai_implementation", self.default_models["openai"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") + for attempt in range(max_retries): try: # Try max_tokens first, fallback to max_completion_tokens if unsupported try: response = await client.chat.completions.create( - model=self.default_models["openai"], + model=impl_model, messages=openai_messages, tools=openai_tools if openai_tools else None, max_tokens=max_tokens, @@ -1023,7 +1032,7 @@ async def _call_openai_with_tools( if "max_tokens" in str(e) and "max_completion_tokens" in str(e): # Retry with max_completion_tokens for models that require it response = await client.chat.completions.create( - model=self.default_models["openai"], + model=impl_model, messages=openai_messages, tools=openai_tools if openai_tools else None, max_completion_tokens=max_tokens, diff --git a/workflows/code_implementation_workflow_index.py b/workflows/code_implementation_workflow_index.py index 079754d9..8ddf00ab 100644 --- a/workflows/code_implementation_workflow_index.py +++ b/workflows/code_implementation_workflow_index.py @@ -33,7 +33,7 @@ from workflows.agents import CodeImplementationAgent from workflows.agents.memory_agent_concise import ConciseMemoryAgent from config.mcp_tool_definitions_index import get_mcp_tools -from utils.llm_utils import get_preferred_llm_class, get_default_models +from utils.llm_utils import get_preferred_llm_class, get_default_models, load_api_config # DialogueLogger removed - no longer needed @@ -62,10 +62,9 @@ def __init__(self, config_path: str = "mcp_agent.secrets.yaml"): ) def _load_api_config(self) -> Dict[str, Any]: - """Load API configuration from YAML file""" + """Load API configuration with environment variable override.""" try: - with open(self.config_path, "r", encoding="utf-8") as f: - return yaml.safe_load(f) + return load_api_config(self.config_path) except Exception as e: raise Exception(f"Failed to load API config: {e}") @@ -672,8 +671,11 @@ async def _call_anthropic_with_tools( ] try: + # Use implementation-specific model for code generation + impl_model = self.default_models.get("anthropic_implementation", self.default_models["anthropic"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") response = await client.messages.create( - model=self.default_models["anthropic"], + model=impl_model, system=system_message, messages=validated_messages, tools=tools, @@ -774,8 +776,11 @@ async def _call_google_with_tools( try: # Google Gemini API call using the native SDK # client is google.genai.Client instance + # Use implementation-specific model for code generation + impl_model = self.default_models.get("google_implementation", self.default_models["google"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") response = await client.aio.models.generate_content( - model=self.default_models["google"], + model=impl_model, contents=gemini_messages, config=config, ) @@ -998,12 +1003,16 @@ async def _call_openai_with_tools( max_retries = 3 retry_delay = 2 # seconds + # Use implementation-specific model for code generation + impl_model = self.default_models.get("openai_implementation", self.default_models["openai"]) + self.logger.info(f"πŸ”§ Code generation using model: {impl_model}") + for attempt in range(max_retries): try: # Try max_tokens first, fallback to max_completion_tokens if unsupported try: response = await client.chat.completions.create( - model=self.default_models["openai"], + model=impl_model, messages=openai_messages, tools=openai_tools if openai_tools else None, max_tokens=max_tokens, @@ -1013,7 +1022,7 @@ async def _call_openai_with_tools( if "max_tokens" in str(e) and "max_completion_tokens" in str(e): # Retry with max_completion_tokens for models that require it response = await client.chat.completions.create( - model=self.default_models["openai"], + model=impl_model, messages=openai_messages, tools=openai_tools if openai_tools else None, max_completion_tokens=max_tokens,