[FEAT][Implement MCP into agent.py]

.gitignore

@@ -404,4 +404,3 @@ flycheck_*.el
 
 # network security
 /network-security.data
-

agent_exec_benchmark.py

@@ -0,0 +1,284 @@
+import asyncio
+import concurrent.futures
+import json
+import os
+import psutil
+import datetime
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+from swarms.structs.agent import Agent
+from loguru import logger
+
+
+class AgentBenchmark:
+    def __init__(
+        self,
+        num_iterations: int = 5,
+        output_dir: str = "benchmark_results",
+    ):
+        self.num_iterations = num_iterations
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(exist_ok=True)
+
+        # Use process pool for CPU-bound tasks
+        self.process_pool = concurrent.futures.ProcessPoolExecutor(
+            max_workers=min(os.cpu_count(), 4)
+        )
+
+        # Use thread pool for I/O-bound tasks
+        self.thread_pool = concurrent.futures.ThreadPoolExecutor(
+            max_workers=min(os.cpu_count() * 2, 8)
+        )
+
+        self.default_queries = [
+            "Conduct an analysis of the best real undervalued ETFs",
+            "What are the top performing tech stocks this quarter?",
+            "Analyze current market trends in renewable energy sector",
+            "Compare Bitcoin and Ethereum investment potential",
+            "Evaluate the risk factors in emerging markets",
+        ]
+
+        self.agent = self._initialize_agent()
+        self.process = psutil.Process()
+
+        # Cache for storing repeated query results
+        self._query_cache = {}
+
+    def _initialize_agent(self) -> Agent:
+        return Agent(
+            agent_name="Financial-Analysis-Agent",
+            agent_description="Personal finance advisor agent",
+            # system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
+            max_loops=1,
+            model_name="gpt-4o-mini",
+            dynamic_temperature_enabled=True,
+            interactive=False,
+        )
+
+    def _get_system_metrics(self) -> Dict[str, float]:
+        # Optimized system metrics collection
+        return {
+            "cpu_percent": self.process.cpu_percent(),
+            "memory_mb": self.process.memory_info().rss / 1024 / 1024,
+        }
+
+    def _calculate_statistics(
+        self, values: List[float]
+    ) -> Dict[str, float]:
+        if not values:
+            return {}
+
+        sorted_values = sorted(values)
+        n = len(sorted_values)
+        mean_val = sum(values) / n
+
+        stats = {
+            "mean": mean_val,
+            "median": sorted_values[n // 2],
+            "min": sorted_values[0],
+            "max": sorted_values[-1],
+        }
+
+        # Only calculate stdev if we have enough values
+        if n > 1:
+            stats["std_dev"] = (
+                sum((x - mean_val) ** 2 for x in values) / n
+            ) ** 0.5
+
+        return {k: round(v, 3) for k, v in stats.items()}
+
+    async def process_iteration(
+        self, query: str, iteration: int
+    ) -> Dict[str, Any]:
+        """Process a single iteration of a query"""
+        try:
+            # Check cache for repeated queries
+            cache_key = f"{query}_{iteration}"
+            if cache_key in self._query_cache:
+                return self._query_cache[cache_key]
+
+            iteration_start = datetime.datetime.now()
+            pre_metrics = self._get_system_metrics()
+
+            # Run the agent
+            try:
+                self.agent.run(query)
+                success = True
+            except Exception as e:
+                str(e)
+                success = False
+
+            execution_time = (
+                datetime.datetime.now() - iteration_start
+            ).total_seconds()
+            post_metrics = self._get_system_metrics()
+
+            result = {
+                "execution_time": execution_time,
+                "success": success,
+                "pre_metrics": pre_metrics,
+                "post_metrics": post_metrics,
+                "iteration_data": {
+                    "iteration": iteration + 1,
+                    "execution_time": round(execution_time, 3),
+                    "success": success,
+                    "system_metrics": {
+                        "pre": pre_metrics,
+                        "post": post_metrics,
+                    },
+                },
+            }
+
+            # Cache the result
+            self._query_cache[cache_key] = result
+            return result
+
+        except Exception as e:
+            logger.error(f"Error in iteration {iteration}: {e}")
+            raise
+
+    async def run_benchmark(
+        self, queries: Optional[List[str]] = None
+    ) -> Dict[str, Any]:
+        """Run the benchmark asynchronously"""
+        queries = queries or self.default_queries
+        benchmark_data = {
+            "metadata": {
+                "timestamp": datetime.datetime.now().isoformat(),
+                "num_iterations": self.num_iterations,
+                "agent_config": {
+                    "model_name": self.agent.model_name,
+                    "max_loops": self.agent.max_loops,
+                },
+            },
+            "results": {},
+        }
+
+        async def process_query(query: str):
+            query_results = {
+                "execution_times": [],
+                "system_metrics": [],
+                "iterations": [],
+            }
+
+            # Process iterations concurrently
+            tasks = [
+                self.process_iteration(query, i)
+                for i in range(self.num_iterations)
+            ]
+            iteration_results = await asyncio.gather(*tasks)
+
+            for result in iteration_results:
+                query_results["execution_times"].append(
+                    result["execution_time"]
+                )
+                query_results["system_metrics"].append(
+                    result["post_metrics"]
+                )
+                query_results["iterations"].append(
+                    result["iteration_data"]
+                )
+
+            # Calculate statistics
+            query_results["statistics"] = {
+                "execution_time": self._calculate_statistics(
+                    query_results["execution_times"]
+                ),
+                "memory_usage": self._calculate_statistics(
+                    [
+                        m["memory_mb"]
+                        for m in query_results["system_metrics"]
+                    ]
+                ),
+                "cpu_usage": self._calculate_statistics(
+                    [
+                        m["cpu_percent"]
+                        for m in query_results["system_metrics"]
+                    ]
+                ),
+            }
+
+            return query, query_results
+
+        # Execute all queries concurrently
+        query_tasks = [process_query(query) for query in queries]
+        query_results = await asyncio.gather(*query_tasks)
+
+        for query, results in query_results:
+            benchmark_data["results"][query] = results
+
+        return benchmark_data
+
+    def save_results(self, benchmark_data: Dict[str, Any]) -> str:
+        """Save benchmark results efficiently"""
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = (
+            self.output_dir / f"benchmark_results_{timestamp}.json"
+        )
+
+        # Write results in a single operation
+        with open(filename, "w") as f:
+            json.dump(benchmark_data, f, indent=2)
+
+        logger.info(f"Benchmark results saved to: {filename}")
+        return str(filename)
+
+    def print_summary(self, results: Dict[str, Any]):
+        """Print a summary of the benchmark results"""
+        print("\n=== Benchmark Summary ===")
+        for query, data in results["results"].items():
+            print(f"\nQuery: {query[:50]}...")
+            stats = data["statistics"]["execution_time"]
+            print(f"Average time: {stats['mean']:.2f}s")
+            print(
+                f"Memory usage (avg): {data['statistics']['memory_usage']['mean']:.1f}MB"
+            )
+            print(
+                f"CPU usage (avg): {data['statistics']['cpu_usage']['mean']:.1f}%"
+            )
+
+    async def run_with_timeout(
+        self, timeout: int = 300
+    ) -> Dict[str, Any]:
+        """Run benchmark with timeout"""
+        try:
+            return await asyncio.wait_for(
+                self.run_benchmark(), timeout
+            )
+        except asyncio.TimeoutError:
+            logger.error(
+                f"Benchmark timed out after {timeout} seconds"
+            )
+            raise
+
+    def cleanup(self):
+        """Cleanup resources"""
+        self.process_pool.shutdown()
+        self.thread_pool.shutdown()
+        self._query_cache.clear()
+
+
+async def main():
+    try:
+        # Create and run benchmark
+        benchmark = AgentBenchmark(num_iterations=1)
+
+        # Run benchmark with timeout
+        results = await benchmark.run_with_timeout(timeout=300)
+
+        # Save results
+        benchmark.save_results(results)
+
+        # Print summary
+        benchmark.print_summary(results)
+
+    except Exception as e:
+        logger.error(f"Benchmark failed: {e}")
+    finally:
+        # Cleanup resources
+        benchmark.cleanup()
+
+
+if __name__ == "__main__":
+    # Run the async main function
+    asyncio.run(main())