Merge pull request #83 from zhuyanhuazhuyanhua/main

feat：add thread pool

Author: Takki-cd

oxygent/oxy/function_tools/function_hub.py

@@ -7,6 +7,7 @@
 
 import asyncio
 import functools
+import concurrent.futures
 
 from pydantic import Field
 
@@ -29,6 +30,18 @@ class FunctionHub(BaseTool):
         default_factory=dict, description="Registry of functions and their metadata"
     )
 
+    def __init__(self, **data):
+        """Initialize the FunctionHub with thread pool support."""
+        super().__init__(**data)
+        self._thread_pool = None  # Private attribute for thread pool
+
+    @property
+    def thread_pool(self):
+        """Lazy initialization of thread pool."""
+        if self._thread_pool is None:
+            self._thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=4)
+        return self._thread_pool
+
     async def init(self):
         """Initialize the hub by creating FunctionTool instances for all registered
         functions.
@@ -67,14 +80,35 @@ def decorator(func):
             if asyncio.iscoroutinefunction(func):
                 async_func = func
             else:
-                # Wrap synchronous function to make it asynchronous
+                # Wrap synchronous function to make it asynchronous using thread pool
                 @functools.wraps(func)
                 async def async_func(*args, **kwargs):
-                    # TODO: Use thread pool for blocking synchronous operations
-                    return func(*args, **kwargs)
+                    # Use thread pool for blocking synchronous operations
+                    loop = asyncio.get_event_loop()
+                    if kwargs:
+                        # 如果有kwargs，使用functools.partial包装函数
+                        partial_func = functools.partial(func, **kwargs)
+                        return await loop.run_in_executor(
+                            self.thread_pool,
+                            partial_func,
+                            *args
+                        )
+                    else:
+                        # 如果没有kwargs，直接调用
+                        return await loop.run_in_executor(
+                            self.thread_pool,
+                            func,
+                            *args
+                        )
 
             # Register function in the hub's dictionary
             self.func_dict[func.__name__] = (description, async_func)
             return async_func  # Return the async version
 
         return decorator
+
+    async def cleanup(self):
+        """Clean up resources, including the thread pool."""
+        if self._thread_pool:
+            self._thread_pool.shutdown(wait=True)
+            self._thread_pool = None

test/unittest/test_function_hub.py

@@ -3,6 +3,8 @@
 """
 
 import asyncio
+import time
+import concurrent.futures
 
 import pytest
 
@@ -92,3 +94,140 @@ def inc(x: int):
 
     result = asyncio.run(async_inc(41))
     assert result == 42
+
+
+# ────────────────────────────────────────────────────────────────────────────
+# Thread Pool Tests
+# ────────────────────────────────────────────────────────────────────────────
+def test_thread_pool_lazy_initialization(func_hub):
+    """Test that thread pool is lazily initialized."""
+    # Thread pool should be None initially
+    assert func_hub._thread_pool is None
+    
+    # Accessing thread_pool property should initialize it
+    pool = func_hub.thread_pool
+    assert isinstance(pool, concurrent.futures.ThreadPoolExecutor)
+    assert func_hub._thread_pool is not None
+    
+    # Second access should return the same instance
+    pool2 = func_hub.thread_pool
+    assert pool is pool2
+
+
+@pytest.mark.asyncio
+async def test_sync_function_execution_with_thread_pool(func_hub):
+    """Test that sync functions are executed in thread pool."""
+    execution_info = {"thread_id": None, "main_thread_id": None}
+    
+    @func_hub.tool("test sync function")
+    def blocking_function(duration: float):
+        """Simulate blocking operation."""
+        import threading
+        execution_info["thread_id"] = threading.current_thread().ident
+        time.sleep(duration)
+        return f"completed in thread {execution_info['thread_id']}"
+    
+    # Get main thread ID
+    import threading
+    execution_info["main_thread_id"] = threading.current_thread().ident
+    
+    # Execute the function
+    _, async_func = func_hub.func_dict["blocking_function"]
+    result = await async_func(0.1)
+    
+    # Verify function executed in different thread
+    assert execution_info["thread_id"] is not None
+    assert execution_info["thread_id"] != execution_info["main_thread_id"]
+    assert "completed in thread" in result
+
+
+@pytest.mark.asyncio
+async def test_sync_function_with_kwargs_in_thread_pool(func_hub):
+    """Test sync functions with kwargs are executed in thread pool."""
+    @func_hub.tool("test function with kwargs")
+    def function_with_kwargs(a: int, b: int, multiplier: float = 1.0):
+        """Function that uses kwargs."""
+        time.sleep(0.01)  # Small delay to simulate work
+        return (a + b) * multiplier
+    
+    # Execute function with kwargs
+    _, async_func = func_hub.func_dict["function_with_kwargs"]
+    result = await async_func(2, 3, multiplier=2.0)
+    
+    assert result == 10.0  # (2 + 3) * 2.0 = 10.0
+
+
+@pytest.mark.asyncio
+async def test_cleanup_shuts_down_thread_pool(func_hub):
+    """Test that cleanup properly shuts down thread pool."""
+    # Initialize thread pool by accessing it
+    pool = func_hub.thread_pool
+    assert isinstance(pool, concurrent.futures.ThreadPoolExecutor)
+    
+    # Verify pool is active
+    assert func_hub._thread_pool is not None
+    
+    # Cleanup should shut down the thread pool
+    await func_hub.cleanup()
+    
+    # Thread pool should be None after cleanup
+    assert func_hub._thread_pool is None
+
+
+@pytest.mark.asyncio
+async def test_multiple_cleanup_calls_safe(func_hub):
+    """Test that multiple cleanup calls are safe."""
+    # Initialize thread pool
+    func_hub.thread_pool
+    
+    # First cleanup
+    await func_hub.cleanup()
+    assert func_hub._thread_pool is None
+    
+    # Second cleanup should not raise error
+    await func_hub.cleanup()
+    assert func_hub._thread_pool is None
+
+
+@pytest.mark.asyncio
+async def test_cleanup_without_thread_pool_initialization(func_hub):
+    """Test cleanup when thread pool was never initialized."""
+    # Ensure thread pool is not initialized
+    assert func_hub._thread_pool is None
+    
+    # Cleanup should work without errors
+    await func_hub.cleanup()
+    assert func_hub._thread_pool is None
+
+
+@pytest.mark.asyncio
+async def test_concurrent_sync_function_execution(func_hub):
+    """Test concurrent execution of multiple sync functions."""
+    results = []
+    
+    @func_hub.tool("concurrent task")
+    def concurrent_task(task_id: int, duration: float):
+        """Simulate concurrent blocking operation."""
+        time.sleep(duration)
+        return f"task_{task_id}_completed"
+    
+    # Execute multiple tasks concurrently
+    _, async_func = func_hub.func_dict["concurrent_task"]
+    tasks = [
+        async_func(1, 0.1),
+        async_func(2, 0.15),
+        async_func(3, 0.05),
+    ]
+    
+    start_time = time.time()
+    results = await asyncio.gather(*tasks)
+    total_time = time.time() - start_time
+    
+    # Verify all tasks completed
+    assert len(results) == 3
+    assert "task_1_completed" in results
+    assert "task_2_completed" in results
+    assert "task_3_completed" in results
+    
+    # Total time should be less than sum of individual times (due to concurrency)
+    assert total_time < 0.3  # Should be around 0.15s (max duration)