support concurrency

lang_agent/dummy/calculator.py

@@ -1,5 +1,6 @@
 import math
 import random
+import asyncio
 from dataclasses import dataclass, field
 from typing import Type, List
 import tyro
@@ -8,6 +9,19 @@ import time
 from lang_agent.config import ToolConfig
 from lang_agent.base import LangToolBase
 
+# Concurrency control: limit max concurrent calculations
+MAX_CONCURRENT_CALCULATIONS = 10
+_semaphore: asyncio.Semaphore | None = None
+
+
+def get_semaphore() -> asyncio.Semaphore:
+    """Lazy initialization of semaphore (must be created within an event loop)."""
+    global _semaphore
+    if _semaphore is None:
+        _semaphore = asyncio.Semaphore(MAX_CONCURRENT_CALCULATIONS)
+    return _semaphore
+
+
 @tyro.conf.configure(tyro.conf.SuppressFixed)
 @dataclass
 class CalculatorConfig(ToolConfig):
@@ -20,10 +34,18 @@ class Calculator(LangToolBase):
 
     def calculator(self, python_expression: str) -> dict:
         """For mathamatical calculation, always use this tool to calculate the result of a python expression. You can use 'math' or 'random' directly, without 'import'."""
-        # time.sleep(2)
         result = eval(python_expression, {"math": math, "random": random})
         return {"success": True, "result": result}
 
+    async def calculator_async(self, python_expression: str) -> dict:
+        """Async version: runs eval in a thread pool to avoid blocking the event loop."""
+        async with get_semaphore():
+            await asyncio.sleep(5)  # Simulate delay for testing
+            result = await asyncio.to_thread(
+                eval, python_expression, {"math": math, "random": random}
+            )
+            return {"success": True, "result": result}
+
     def get_tool_fnc(self):
         return [self.calculator]
 
@@ -38,9 +60,9 @@ mcp = FastMCP("Calculator Server")
 
 
 @mcp.tool()
-def calculate(python_expression: str) -> dict:
+async def calculate(python_expression: str) -> dict:
     """For mathematical calculation, always use this tool to calculate the result of a python expression. You can use 'math' or 'random' directly, without 'import'."""
-    return calculator.calculator(python_expression)
+    return await calculator.calculator_async(python_expression)
 
 
 if __name__ == "__main__":