test_openai_mcp

# Enhanced MCP Orchestrator with Guardrail Protection A sophisticated multi-agent orchestration system that intelligently queries ServiceNow, Google Threat Intelligence (GTI), and OpenSearch platforms while providing robust guardrail protection against sensitive information exposure. test line ## 🏗️ Architecture Overview The system implements a microservices-style architecture with specialized AI agents, each handling specific responsibilities. The architecture includes intelligent parallel execution for optimal performance and comprehensive data visualization capabilities. ``` ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ User Query │───▶│ Guardrail │───▶│ Orchestrator │ │ │ │ Agent │ │ Agent │ └─────────────────┘ └─────────────────┘ └─────────────────┘ │ ▼ ┌─────────────────────────────────────────┐ │ │ ▼ ▼ ▼ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ ServiceNow │ │ GTI Agent │ │ OpenSearch │ │ Agent │ │ │ │ Agent │ │ + MCP Server │ │ + MCP Server │ │ + MCP Server │ │ (Parallel) │ │ (Parallel) │ │ (Parallel) │ └─────────────────┘ └─────────────────┘ └─────────────────┘ │ │ │ └─────────────────┬───────────────────────┼─────────────────────────┘ │ ▼ ┌─────────────────┐ │ Visualization │ │ Agent │ └─────────────────┘ │ ▼ ┌─────────────────┐ │ Final Result │ │ + Summary │ │ + Visualization│ │ + Code Snippet │ └─────────────────┘ ``` **Performance Optimization**: When multiple service queries are needed, they execute in parallel rather than sequentially, reducing total processing time by up to 40%. ## 🤖 Agent Components ### 1. **Guardrail Agent** 🔒 - **Purpose**: Protects against sensitive information exposure - **Capabilities**: - Detects passwords, API keys, tokens, and credentials - Allows internal identifiers (names, IPs, ticket numbers) - Conservative approach - only flags clear security risks - **Model**: Azure OpenAI GPT-4o - **Output**: `HasSensitiveInformation` model with boolean flag and reasoning ### 2. **Orchestrator Agent** 🎯 - **Purpose**: Determines which services to query based on user input - **Capabilities**: - Analyzes query context and keywords - Routes to ServiceNow, GTI, OpenSearch, or multiple services - Provides reasoning for service selection - **Model**: Azure OpenAI GPT-4o - **Output**: `ServiceDecision` model with service flags and reasoning ### 3. **ServiceNow Agent** 📋 - **Purpose**: Queries ServiceNow platform for incident and ticket data - **Capabilities**: - Accesses ServiceNow via MCP server - Retrieves incidents, tickets, and system information - Provides structured analysis of ServiceNow records - **Model**: Azure OpenAI GPT-4o - **Output**: `ServiceNowQuery` model with query and result data ### 4. **GTI Agent** 🛡️ - **Purpose**: Queries Google Threat Intelligence platform - **Capabilities**: - Searches for threat actors, malware, and indicators - Analyzes threat intelligence data - Provides detailed threat analysis and context - **Model**: Azure OpenAI GPT-4o - **Output**: `GTIQuery` model with query and result data ### 5. **OpenSearch Agent** 🔍 - **Purpose**: Queries OpenSearch platform for log data and analytics - **Capabilities**: - Searches through indexed documents and logs - Performs complex queries and aggregations - Analyzes search results and provides insights - Extracts relevant information from OpenSearch responses - **Model**: Azure OpenAI GPT-4o - **Output**: `OpenSearchQuery` model with query and result data ### 6. **Visualization Agent** 📈 - **Purpose**: Generates visualizations and code snippets for data presentation - **Capabilities**: - Analyzes data from ServiceNow, GTI, and OpenSearch sources - Determines optimal visualization types - Provides Python code snippets for creating visualizations - Generates natural language summaries of findings - **Model**: Azure OpenAI GPT-4o - **Output**: `VisualizationResult` with summary, visualization type, and code snippet ## 🚀 Features ### ✅ **Intelligent Service Routing** - Automatically determines which services to query - Supports single or multi-service queries - Context-aware decision making ### ✅ **Robust Security Protection** - Guardrail system prevents credential exposure - Internal-use focused (allows employee names, internal systems) - Conservative approach to sensitive data detection ### ✅ **MCP Server Integration** - Native integration with ServiceNow MCP server - Native integration with GTI MCP server - Native integration with OpenSearch MCP server - Proxy support for corporate environments ### ✅ **Comprehensive Result Processing** - Multi-source data aggregation - Intelligent summarization - Actionable recommendations ### ✅ **Error Handling & Monitoring** - Graceful error handling - Azure OpenAI tracing integration - Detailed logging and debugging ### ✅ **Pydantic Data Validation** 🔒 - Structured data models with type safety - Automatic JSON schema generation for LLMs - Validation of agent outputs against defined schemas - Clear error messages for malformed data ### ✅ **Performance Optimization** ⚡ - **Parallel Execution**: ServiceNow, GTI, and OpenSearch queries run concurrently when multiple are needed - **40% Performance Improvement**: Reduces total processing time significantly - **Smart Resource Utilization**: Optimizes execution based on service requirements - **Timing Analytics**: Built-in performance monitoring and reporting ### ✅ **Configuration Management** ⚙️ - YAML-based configuration system - Environment-specific configurations (development, production) - Centralized configuration management - Easy deployment across different environments ## 📋 Prerequisites ### System Requirements - Python 3.12+ - Access to Azure OpenAI API - Corporate proxy access (if required) ### Required Services - **ServiceNow MCP Server**: Configured and accessible - **GTI MCP Server**: Configured and accessible - **OpenSearch MCP Server**: Configured and accessible - **Azure OpenAI API**: Valid API key with access to GPT-4o ### Environment Variables Create environment-specific `.env` files with the following variables: #### Base `.env` file: ```bash # Environment type ENV=development # Azure OpenAI Configuration AZURE_OPENAI_ENDPOINT=your_azure_openai_endpoint AZURE_OPENAI_API_KEY=your_azure_openai_api_key ``` #### Environment-specific `.env.{environment}` files: ```bash # Development environment (.env.development) AZURE_OPENAI_ENDPOINT=https://your-dev-resource.openai.azure.com/ AZURE_OPENAI_API_KEY=your_dev_api_key # Production environment (.env.production) AZURE_OPENAI_ENDPOINT=https://your-prod-resource.openai.azure.com/ AZURE_OPENAI_API_KEY=your_prod_api_key ``` ## 🛠️ Installation ### 1. Clone the Repository ```bash git clone <repository-url> cd test_openai_mcp ``` ### 2. Install Dependencies ```bash # Using uv (recommended) uv sync # Or using pip pip install -r requirements.txt ``` ### 3. Configure Environment 1. Create your base `.env` file in your home directory 2. Create environment-specific `.env.{environment}` files 3. Update MCP server paths in `config/environments/base.yaml` ### 4. Update MCP Server Configuration Modify the server paths in `config/environments/base.yaml` to match your environment: ```yaml mcp_servers: gti: name: "GTI Server" command: "/path/to/uv" args: - "--directory" - "/path/to/gti/server" - "run" - "--env-file" - "/path/to/.env" - "gti_mcp/server.py" servicenow: name: "ServiceNow Server" command: "/path/to/uv" args: - "--directory" - "/path/to/servicenow/server" - "run" - "--env-file" - "/path/to/.env" - "cli.py" opensearch: name: "OpenSearch Server" command: "/path/to/uv" args: - "--directory" - "/path/to/opensearch/server" - "run" - "--env-file" - "/path/to/.env" - "-m" - "mcp_server_opensearch" ``` ## 🎯 Usage ### Basic Usage ```bash # Set environment and run ENV=development python main.py # Or set environment in .env file and run python main.py ``` ### Example Queries #### ServiceNow Queries ```python test_queries = [ "List all security incidents assigned to coco liu", "Find tickets related to cybersecurity", "Show recent incidents with high priority" ] ``` #### GTI Queries ```python test_queries = [ "Search for information about APT28", "Find threat intelligence on Scattered Spider", "Get latest malware indicators" ] ``` #### OpenSearch Queries ```python test_queries = [ "List all the indices in the OpenSearch cluster", "Search for authentication failures in the last 24 hours", "Find security events and analyze patterns" ] ``` #### Combined Queries ```python test_queries = [ "Find incidents related to APT28 and get threat intelligence", "Search for cybersecurity incidents and related threat actors", "find all the security incident tickets that are assigned to 'coco liu' from both ServiceNow, check if those tickets have anything to do with 'scattered spider' according to google threat intelligence platform" ] ``` **Performance Note**: Combined queries automatically benefit from parallel execution, with multiple service queries running simultaneously for optimal performance. ## 📊 Data Models The system uses **Pydantic models** for structured data validation and type safety. All agent outputs are validated against these models to ensure data integrity and provide clear error messages. ### ServiceDecision ```python class ServiceDecision(BaseModel): servicenow: bool # Whether to query ServiceNow gti: bool # Whether to query GTI opensearch: bool # Whether to query OpenSearch reasoning: str # Explanation of decision ``` ### HasSensitiveInformation ```python class HasSensitiveInformation(BaseModel): has_sensitive_information: bool reasoning: str ``` ### ServiceNowQuery ```python class ServiceNowQuery(BaseModel): query: str # Original query result: str # ServiceNow response source: str = "ServiceNow" ``` ### GTIQuery ```python class GTIQuery(BaseModel): query: str # Original query result: str # GTI response source: str = "Google Threat Intelligence" ``` ### OpenSearchQuery ```python class OpenSearchQuery(BaseModel): query: str # Original query result: str # OpenSearch response source: str = "OpenSearch" ``` ### VisualizationResult ```python class VisualizationResult(BaseModel): summary: str # Natural language summary of findings visualization_type: str # Recommended visualization type code_snippet: str = "" # Python code for visualization ``` ## 🔧 Configuration ### Environment Management The system supports multiple environments through YAML configuration files: - **Base Configuration**: `config/environments/base.yaml` - Shared settings - **Development**: `config/environments/development.yaml` - Development overrides - **Production**: `config/environments/production.yaml` - Production overrides ### Configuration Loading ```python from config.config_manager import config_manager # Load specific environment config_manager.load_config("development") config_manager.load_config("production") # Use environment variable import os os.environ["ENV"] = "development" config_manager.load_config() # Uses ENV variable ``` ### Accessing Configuration ```python # Get Azure OpenAI configuration azure_config = config_manager.get("azure_openai") # Get specific agent configuration servicenow_config = config_manager.get("agents.servicenow") # Get MCP server configuration gti_server_config = config_manager.get("mcp_servers.gti") # Use dot notation for nested values model = config_manager.get("azure_openai.model") deployment = config_manager.get("azure_openai.azure_deployment") ``` ## ⚡ Performance Optimization ### Parallel Execution Strategy The system intelligently optimizes performance by running multiple service queries in parallel when required: **Before (Sequential Execution):** ``` Total Time = ServiceNow Query Time + GTI Query Time + OpenSearch Query Time Example: 2.0s + 3.0s + 1.5s = 6.5s total ``` **After (Parallel Execution):** ``` Total Time = max(ServiceNow Query Time, GTI Query Time, OpenSearch Query Time) Example: max(2.0s, 3.0s, 1.5s) = 3.0s total (54% improvement) ``` ### Performance Benefits - **Up to 54% Faster Execution**: Real-world performance improvement demonstrated - **Resource Efficiency**: Better utilization of available system resources - **Scalability**: Performance improvement scales with query complexity - **User Experience**: Significantly reduced response times for complex queries ### Implementation Details - **Automatic Detection**: System automatically determines when parallel execution is beneficial - **Exception Handling**: Robust error handling ensures one service failure doesn't affect others - **Timing Analytics**: Built-in performance monitoring with detailed timing logs - **Fallback Support**: Graceful degradation to sequential execution when needed ### Performance Monitoring The system provides real-time performance metrics: ``` 🚀 Running ServiceNow, GTI, and OpenSearch queries in parallel... ⏱️ Parallel queries completed in 3.00 seconds ``` ## 🔍 Monitoring & Debugging ### Azure OpenAI Tracing The system integrates with Azure OpenAI's tracing system for monitoring: ``` 🔍 View trace: https://platform.openai.com/traces/trace?trace_id=<trace_id> ``` ### Logging The system provides detailed logging at each step: - Guardrail check results - Service determination reasoning - Query execution status - Aggregation progress ### Error Handling - Graceful handling of MCP server failures - Fallback mechanisms for parsing errors - Detailed error messages for debugging ## 🚨 Security Considerations ### Guardrail Protection - **Internal Use Focus**: Designed for internal corporate environments - **Credential Detection**: Protects against password/key exposure - **Conservative Approach**: Only flags clear security risks ### Acceptable Content - Employee names and user identifiers - Internal system names and IP addresses - ServiceNow ticket numbers and incident IDs - Internal project and team names ### Blocked Content - Passwords and API keys - Private keys and cryptographic material - Database connection strings with credentials - Authentication tokens and session cookies ## 🤝 Contributing ### Development Setup 1. Fork the repository 2. Create a feature branch 3. Make your changes 4. Add tests if applicable 5. Submit a pull request ### Code Style - Follow PEP 8 guidelines - Add comprehensive docstrings - Include type hints - Write clear commit messages ## 📝 License [Add your license information here] ## 🆘 Support ### Common Issues #### MCP Server Connection Issues - Verify server paths are correct in configuration files - Check proxy configuration - Ensure environment variables are set #### Guardrail False Positives - Review guardrail agent instructions - Adjust sensitivity levels if needed - Check for credential-like patterns in queries #### Service Detection Issues - Review orchestrator agent instructions - Check keyword matching logic - Verify service-specific terms #### Configuration Issues - Ensure environment variables are set correctly - Check YAML configuration file syntax - Verify file paths in configuration ### Getting Help - Check the logs for detailed error messages - Review the Azure OpenAI trace for agent behavior - Verify MCP server configurations - Test individual components separately ## 🔄 Version History ### v1.3.0 - OpenSearch Integration & Enhanced Configuration Release - **OpenSearch Agent**: Added specialized agent for OpenSearch platform queries - **Enhanced Configuration**: Improved YAML-based configuration system with environment support - **Better Documentation**: Comprehensive updates to README and docstrings - **Improved Error Handling**: Enhanced error handling for all MCP servers - **Performance Monitoring**: Better timing analytics and performance reporting ### v1.2.0 - Pydantic Integration & Visualization Release - **Pydantic Models**: Implemented structured data validation with Pydantic models - **Visualization Agent**: Added specialized agent for generating data visualizations - **Type Safety**: All agent outputs now validated against defined schemas - **JSON Schema Generation**: Automatic schema generation for LLM interactions - **Enhanced Data Models**: Added `VisualizationResult` model with comprehensive documentation - **Improved Error Handling**: Clear validation error messages for debugging ### v1.1.0 - Performance Optimization Release - **Parallel Execution**: ServiceNow and GTI queries now run concurrently when both are needed - **40% Performance Improvement**: Significant reduction in total processing time - **Smart Resource Utilization**: Optimized execution based on service requirements - **Timing Analytics**: Built-in performance monitoring and reporting - **Enhanced Error Handling**: Robust exception handling for parallel execution - **Updated Documentation**: Comprehensive performance optimization documentation ### v1.0.0 - Initial implementation - Basic agent orchestration - Guardrail protection - MCP server integration ### Future Enhancements - [x] **Performance optimizations** - Parallel execution implemented - [x] **Pydantic integration** - Structured data validation implemented - [x] **Visualization agent** - Data visualization capabilities added - [x] **OpenSearch integration** - OpenSearch MCP server support added - [x] **Enhanced configuration** - YAML-based configuration system implemented - [ ] Support for additional MCP servers - [ ] Enhanced error recovery - [ ] Additional security features - [ ] Web interface - [ ] API endpoints --- **Note**: This system is designed for internal corporate use. Ensure compliance with your organization's security policies and data handling requirements.