- Added MCP client to discover and convert MCP tools to LangChain Tools - Extended configuration manager to support MCP server configurations - Updated CLI to manage MCP server configurations (list, add, remove, use) - Integrated MCP tools into the agent's tool loading system - Added documentation for MCP integration design
453 lines
18 KiB
Python
453 lines
18 KiB
Python
"""
|
|
Configuration manager for LangCode Agent
|
|
Handles persistent storage of provider configurations with encryption and validation
|
|
"""
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional, List, Tuple
|
|
import copy
|
|
from cryptography.fernet import Fernet
|
|
import base64
|
|
|
|
try:
|
|
import requests
|
|
except ImportError:
|
|
requests = None
|
|
|
|
|
|
class ConfigManager:
|
|
"""Manages LangCode Agent configuration with encryption and validation"""
|
|
|
|
def __init__(self, config_dir: Optional[str] = None):
|
|
"""
|
|
Initialize the config manager
|
|
|
|
Args:
|
|
config_dir: Directory to store config files. If None, uses ~/.langcode
|
|
"""
|
|
if config_dir is None:
|
|
self.config_dir = Path.home() / ".langcode"
|
|
else:
|
|
self.config_dir = Path(config_dir)
|
|
|
|
self.config_file = self.config_dir / "config.json"
|
|
self.key_file = self.config_dir / ".key"
|
|
|
|
# Ensure config directory exists with secure permissions
|
|
self._ensure_config_dir()
|
|
|
|
# Load or create encryption key
|
|
self._encryption_key = self._load_or_create_key()
|
|
self._cipher = Fernet(self._encryption_key)
|
|
|
|
# Load configuration
|
|
self.config = self._load_config()
|
|
|
|
def _ensure_config_dir(self):
|
|
"""Ensure the configuration directory exists with secure permissions"""
|
|
self.config_dir.mkdir(parents=True, exist_ok=True)
|
|
# Set directory permissions to owner-only (700)
|
|
os.chmod(self.config_dir, 0o700)
|
|
|
|
def _load_or_create_key(self) -> bytes:
|
|
"""Load existing encryption key or create a new one"""
|
|
if self.key_file.exists():
|
|
try:
|
|
key = self.key_file.read_bytes()
|
|
# Validate key format
|
|
Fernet(key) # This will raise an exception if invalid
|
|
# Set key file permissions to owner-only (600)
|
|
os.chmod(self.key_file, 0o600)
|
|
return key
|
|
except Exception:
|
|
# If key is invalid or corrupted, create a new one
|
|
pass
|
|
|
|
# Generate new key
|
|
key = Fernet.generate_key()
|
|
self.key_file.write_bytes(key)
|
|
# Set key file permissions to owner-only (600)
|
|
os.chmod(self.key_file, 0o600)
|
|
return key
|
|
|
|
def _encrypt_data(self, data: str) -> str:
|
|
"""Encrypt sensitive data"""
|
|
if not data:
|
|
return data
|
|
encrypted = self._cipher.encrypt(data.encode())
|
|
return base64.b64encode(encrypted).decode('utf-8')
|
|
|
|
def _decrypt_data(self, encrypted_data: str) -> str:
|
|
"""Decrypt sensitive data"""
|
|
if not encrypted_data:
|
|
return encrypted_data
|
|
try:
|
|
decoded = base64.b64decode(encrypted_data.encode())
|
|
decrypted = self._cipher.decrypt(decoded)
|
|
return decrypted.decode('utf-8')
|
|
except Exception:
|
|
# If decryption fails, return as-is (might be unencrypted legacy data)
|
|
return encrypted_data
|
|
|
|
def _load_config(self) -> Dict[str, Any]:
|
|
"""Load configuration from file"""
|
|
if self.config_file.exists():
|
|
try:
|
|
with open(self.config_file, 'r') as f:
|
|
config = json.load(f)
|
|
|
|
# Decrypt any encrypted fields in existing config
|
|
return self._decrypt_config(config)
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
# Return default config if file is corrupted
|
|
print(f"Warning: Could not load config file ({e}), using defaults")
|
|
return self._get_default_config()
|
|
else:
|
|
# Return default config if file doesn't exist
|
|
return self._get_default_config()
|
|
|
|
def _decrypt_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Decrypt sensitive fields in configuration"""
|
|
if not isinstance(config, dict):
|
|
return config
|
|
|
|
# Create a deep copy to avoid modifying original
|
|
result = copy.deepcopy(config)
|
|
|
|
# Decrypt configurations
|
|
if "configurations" in result and isinstance(result["configurations"], dict):
|
|
for name, config_data in result["configurations"].items():
|
|
if isinstance(config_data, dict):
|
|
# Decrypt API keys for relevant providers
|
|
provider = config_data.get("provider", "").lower()
|
|
if provider in ["openai", "openrouter", "mcp"] and "api_key" in config_data:
|
|
config_data["api_key"] = self._decrypt_data(config_data["api_key"])
|
|
|
|
return result
|
|
|
|
def _encrypt_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Encrypt sensitive fields in configuration before storage"""
|
|
if not isinstance(config, dict):
|
|
return config
|
|
|
|
# Create a deep copy to avoid modifying original
|
|
result = copy.deepcopy(config)
|
|
|
|
# Encrypt configurations
|
|
if "configurations" in result and isinstance(result["configurations"], dict):
|
|
for name, config_data in result["configurations"].items():
|
|
if isinstance(config_data, dict):
|
|
# Encrypt API keys for relevant providers
|
|
provider = config_data.get("provider", "").lower()
|
|
if provider in ["openai", "openrouter"] and "api_key" in config_data:
|
|
api_key = config_data["api_key"]
|
|
# Only encrypt if not already encrypted (avoid double encryption)
|
|
if api_key and not self._is_likely_encrypted(api_key):
|
|
config_data["api_key"] = self._encrypt_data(api_key)
|
|
|
|
return result
|
|
|
|
def _is_likely_encrypted(self, data: str) -> bool:
|
|
"""Heuristic to check if data appears to be encrypted"""
|
|
if not data or len(data) < 10:
|
|
return False
|
|
# Encrypted data with Fernet is base64 encoded and typically longer
|
|
try:
|
|
# Try to decode as base64
|
|
base64.b64decode(data.encode(), validate=True)
|
|
# If it decodes successfully and is reasonably long, likely encrypted
|
|
return len(data) >= 20
|
|
except Exception:
|
|
return False
|
|
|
|
def _get_default_config(self) -> Dict[str, Any]:
|
|
"""Get default configuration structure"""
|
|
return {
|
|
"configurations": {},
|
|
"active_configuration": None,
|
|
"version": "1.0"
|
|
}
|
|
|
|
def _save_config(self):
|
|
"""Save configuration to file with encryption of sensitive data"""
|
|
try:
|
|
# Encrypt sensitive data before saving
|
|
config_to_save = self._encrypt_config(self.config)
|
|
|
|
# Create a backup of the existing config
|
|
if self.config_file.exists():
|
|
backup_file = self.config_file.with_suffix('.json.bak')
|
|
backup_file.write_text(self.config_file.read_text())
|
|
# Set backup file permissions to owner-only (600)
|
|
os.chmod(backup_file, 0o600)
|
|
|
|
# Write the encrypted config
|
|
with open(self.config_file, 'w') as f:
|
|
json.dump(config_to_save, f, indent=2, sort_keys=True)
|
|
|
|
# Set config file permissions to owner-only (600)
|
|
os.chmod(self.config_file, 0o600)
|
|
except IOError as e:
|
|
raise Exception(f"Failed to save configuration: {e}")
|
|
|
|
def validate_configuration(self, config: Dict[str, Any]) -> Tuple[bool, List[str]]:
|
|
"""
|
|
Validate a configuration dictionary
|
|
|
|
Returns:
|
|
Tuple of (is_valid, list_of_error_messages)
|
|
"""
|
|
errors = []
|
|
|
|
# Check if config is a dictionary
|
|
if not isinstance(config, dict):
|
|
errors.append("Configuration must be a dictionary")
|
|
return False, errors
|
|
|
|
# Validate provider
|
|
provider = config.get("provider", "").lower()
|
|
valid_providers = ["openai", "ollama", "openrouter", "mcp"]
|
|
|
|
if not provider:
|
|
errors.append("Provider is required")
|
|
elif provider not in valid_providers:
|
|
errors.append(f"Unsupported provider: {provider}. Valid options: {', '.join(valid_providers)}")
|
|
|
|
# Validate model name
|
|
model_name = config.get("model_name", "").strip()
|
|
if not model_name:
|
|
errors.append("Model name is required and cannot be empty")
|
|
elif len(model_name) > 100:
|
|
errors.append("Model name is too long (maximum 100 characters)")
|
|
|
|
# Provider-specific validation
|
|
if provider == "openai":
|
|
api_key = config.get("api_key", "").strip()
|
|
if not api_key:
|
|
errors.append("API key is required for OpenAI provider")
|
|
elif not api_key.startswith("sk-"):
|
|
errors.append("OpenAI API key should start with 'sk-'")
|
|
elif len(api_key) < 20:
|
|
errors.append("OpenAI API key appears too short")
|
|
|
|
elif provider == "ollama":
|
|
base_url = config.get("base_url", "").strip()
|
|
if not base_url:
|
|
errors.append("Base URL is required for Ollama provider")
|
|
else:
|
|
# Normalize URL for validation
|
|
if not base_url.startswith(("http://", "https://")):
|
|
errors.append("Ollama base URL must start with http:// or https://")
|
|
else:
|
|
# Remove trailing slash for consistency
|
|
if base_url.endswith("/"):
|
|
config["base_url"] = base_url.rstrip("/")
|
|
|
|
elif provider == "openrouter":
|
|
api_key = config.get("api_key", "").strip()
|
|
base_url = config.get("base_url", "").strip()
|
|
|
|
if not api_key:
|
|
errors.append("API key is required for OpenRouter provider")
|
|
elif len(api_key) < 10:
|
|
errors.append("OpenRouter API key appears too short")
|
|
|
|
if not base_url:
|
|
errors.append("Base URL is required for OpenRouter provider")
|
|
elif not base_url.startswith(("http://", "https://")):
|
|
errors.append("OpenRouter base URL must start with http:// or https://")
|
|
else:
|
|
# Remove trailing slash for consistency
|
|
if base_url.endswith("/"):
|
|
config["base_url"] = base_url.rstrip("/")
|
|
|
|
elif provider == "mcp":
|
|
server_url = config.get("server_url", "").strip()
|
|
if not server_url:
|
|
errors.append("Server URL is required for MCP provider")
|
|
else:
|
|
if not server_url.startswith(("http://", "https://")):
|
|
errors.append("MCP server URL must start with http:// or https://")
|
|
else:
|
|
# Remove trailing slash for consistency
|
|
if server_url.endswith("/"):
|
|
config["server_url"] = server_url.rstrip("/")
|
|
|
|
# API key is optional for MCP (some servers might not require auth)
|
|
api_key = config.get("api_key", "").strip()
|
|
if api_key and len(api_key) < 10:
|
|
errors.append("MCP API key appears too short if provided")
|
|
|
|
# Timeout is optional, but if provided must be positive integer
|
|
timeout = config.get("timeout")
|
|
if timeout is not None:
|
|
try:
|
|
timeout_int = int(timeout)
|
|
if timeout_int <= 0:
|
|
errors.append("MCP timeout must be a positive integer")
|
|
else:
|
|
config["timeout"] = timeout_int
|
|
except ValueError:
|
|
errors.append("MCP timeout must be a valid integer")
|
|
|
|
return len(errors) == 0, errors
|
|
|
|
def get_configurations(self) -> Dict[str, Any]:
|
|
"""Get all configurations"""
|
|
return self.config.get("configurations", {})
|
|
|
|
def get_active_configuration(self) -> Optional[str]:
|
|
"""Get the name of the active configuration"""
|
|
return self.config.get("active_configuration")
|
|
|
|
def get_configuration(self, name: str) -> Optional[Dict[str, Any]]:
|
|
"""Get a specific configuration by name"""
|
|
configurations = self.get_configurations()
|
|
config = configurations.get(name)
|
|
# Return a copy to prevent accidental modification of internal data
|
|
return copy.deepcopy(config) if config else None
|
|
|
|
def add_configuration(self, name: str, config: Dict[str, Any]):
|
|
"""
|
|
Add or update a configuration with validation
|
|
|
|
Args:
|
|
name: Configuration name
|
|
config: Configuration dictionary with provider, model, etc.
|
|
|
|
Raises:
|
|
ValueError: If configuration is invalid
|
|
"""
|
|
if not name or not isinstance(name, str):
|
|
raise ValueError("Configuration name must be a non-empty string")
|
|
|
|
# Validate the configuration
|
|
is_valid, errors = self.validate_configuration(config)
|
|
if not is_valid:
|
|
raise ValueError(f"Invalid configuration: {', '.join(errors)}")
|
|
|
|
if "configurations" not in self.config:
|
|
self.config["configurations"] = {}
|
|
|
|
# Deep copy to avoid reference issues
|
|
self.config["configurations"][name] = copy.deepcopy(config)
|
|
self._save_config()
|
|
|
|
def remove_configuration(self, name: str) -> bool:
|
|
"""
|
|
Remove a configuration
|
|
|
|
Args:
|
|
name: Configuration name to remove
|
|
|
|
Returns:
|
|
True if removed, False if not found
|
|
"""
|
|
if "configurations" in self.config and name in self.config["configurations"]:
|
|
del self.config["configurations"][name]
|
|
|
|
# If this was the active configuration, unset it
|
|
if self.config.get("active_configuration") == name:
|
|
self.config["active_configuration"] = None
|
|
|
|
self._save_config()
|
|
return True
|
|
return False
|
|
|
|
def set_active_configuration(self, name: str) -> bool:
|
|
"""
|
|
Set the active configuration
|
|
|
|
Args:
|
|
name: Configuration name to set as active
|
|
|
|
Returns:
|
|
True if successful, False if configuration not found
|
|
"""
|
|
if name in self.get_configurations():
|
|
self.config["active_configuration"] = name
|
|
self._save_config()
|
|
return True
|
|
return False
|
|
|
|
def get_active_config_details(self) -> Optional[Dict[str, Any]]:
|
|
"""Get the details of the active configuration"""
|
|
active_name = self.get_active_configuration()
|
|
if active_name:
|
|
return self.get_configuration(active_name)
|
|
return None
|
|
|
|
def list_configuration_names(self) -> List[str]:
|
|
"""Get a list of all configuration names"""
|
|
return list(self.get_configurations().keys())
|
|
|
|
def configuration_exists(self, name: str) -> bool:
|
|
"""Check if a configuration exists"""
|
|
return name in self.get_configurations()
|
|
|
|
def test_configuration_connection(self, name: str) -> Tuple[bool, str]:
|
|
"""
|
|
Test connection to a configured provider
|
|
|
|
Args:
|
|
name: Configuration name to test
|
|
|
|
Returns:
|
|
Tuple of (success, message)
|
|
"""
|
|
config = self.get_configuration(name)
|
|
if not config:
|
|
return False, f"Configuration '{name}' not found"
|
|
|
|
provider = config.get("provider", "").lower()
|
|
|
|
# Check if requests is available
|
|
if requests is None:
|
|
return False, "The 'requests' library is not available. Please install it to test connections."
|
|
|
|
try:
|
|
if provider == "ollama":
|
|
base_url = config.get("base_url", "http://localhost:11434")
|
|
# Ensure no trailing slash
|
|
if base_url.endswith("/"):
|
|
base_url = base_url.rstrip("/")
|
|
response = requests.get(f"{base_url}/api/tags", timeout=5)
|
|
if response.status_code == 200:
|
|
models = response.json().get("models", [])
|
|
model_names = [m.get("name", "unknown") for m in models]
|
|
return True, f"Connected successfully. Available models: {', '.join(model_names[:5])}{'...' if len(model_names) > 5 else ''}"
|
|
else:
|
|
return False, f"Ollama server returned status {response.status_code}"
|
|
|
|
elif provider == "openai":
|
|
# For OpenAI, we can't easily test without making an actual API call
|
|
# which might incur costs. Just validate the key format.
|
|
api_key = config.get("api_key", "")
|
|
if api_key and api_key.startswith("sk-") and len(api_key) > 20:
|
|
return True, "API key format appears valid (no live test performed to avoid costs)"
|
|
else:
|
|
return False, "API key format appears invalid"
|
|
|
|
elif provider == "openrouter":
|
|
# Similar to OpenAI, just validate format to avoid costs
|
|
api_key = config.get("api_key", "")
|
|
base_url = config.get("base_url", "https://openrouter.ai/api/v1")
|
|
if api_key and len(api_key) > 10:
|
|
return True, "API key format appears valid (no live test performed to avoid costs)"
|
|
else:
|
|
return False, "API key format appears invalid"
|
|
|
|
else:
|
|
return False, f"Connection testing not implemented for provider: {provider}"
|
|
|
|
except requests.exceptions.Timeout:
|
|
return False, "Connection timeout - please check the URL and network connectivity"
|
|
except requests.exceptions.ConnectionError:
|
|
return False, "Connection failed - please check the URL and ensure the service is running"
|
|
except Exception as e:
|
|
return False, f"Connection test failed: {str(e)}"
|
|
|
|
|
|
# Global config manager instance
|
|
config_manager = ConfigManager() |