From 7c00cf155d210f6e68c35151d66b870c23a5a736 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 16:46:32 +0000
Subject: [PATCH] feat: support both Swarms and OpenAI API keys in integration
 test

- Modified `test_integration.py` to dynamically use either Swarms or OpenAI key
- Enables flexible backend switching for agent API calls
---
 examples/mcp_example/test_integration.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/examples/mcp_example/test_integration.py b/examples/mcp_example/test_integration.py
index 6806feaf..eb33d0df 100644
--- a/examples/mcp_example/test_integration.py
+++ b/examples/mcp_example/test_integration.py
@@ -3,6 +3,15 @@ from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
 from swarms.tools.mcp_integration import MCPServerSseParams
 import logging
 from datetime import datetime
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Use either Swarms API key or OpenAI API key
+api_key = os.getenv("SWARMS_API_KEY") or os.getenv("OPENAI_API_KEY")
+if not api_key:
+    raise ValueError("Please set either SWARMS_API_KEY or OPENAI_API_KEY in your environment")
 
 def setup_agent(name: str, description: str, servers: list) -> Agent:
     """Setup an agent with MCP server connections"""