feat(mcp): cache with one-click toggle (#1269)

* feat: enable cache in mcp

* refactor: remove redundant setting

* fix: conflicts during installation

---------

Co-authored-by: Linlang <Lv.Linlang@hotmail.com>

.gitignore

@@ -178,3 +178,8 @@ rdagent/app/benchmark/factor/example.json
 # UI Server resources
 videos/
 static/
+
+# AI assistant
+.cursor/
+.claude/
+AGENTS.md

rdagent/components/agent/base.py

@@ -1,6 +1,8 @@
 from abc import abstractmethod
 
 import nest_asyncio
+from prefect import task
+from prefect.cache_policies import INPUTS
 from pydantic_ai import Agent
 from pydantic_ai.mcp import MCPServerStreamableHTTP
 
@@ -18,17 +20,50 @@ class BaseAgent:
 
 class PAIAgent(BaseAgent):
     """
-    Pydantic-AI agent
+    Pydantic-AI agent with optional Prefect caching support
     """
 
     agent: Agent
+    enable_cache: bool
 
-    def __init__(self, system_prompt: str, toolsets: list[str | MCPServerStreamableHTTP]):
+    def __init__(
+        self,
+        system_prompt: str,
+        toolsets: list[str | MCPServerStreamableHTTP],
+        enable_cache: bool = False,
+    ):
+        """
+        Initialize Pydantic-AI agent
+
+        Parameters
+        ----------
+        system_prompt : str
+            System prompt for the agent
+        toolsets : list[str | MCPServerStreamableHTTP]
+            List of MCP server URLs or instances
+        enable_cache : bool
+            Enable persistent caching via Prefect. Requires Prefect server:
+            `prefect server start` then set PREFECT_API_URL in environment
+        """
         toolsets = [(ts if isinstance(ts, MCPServerStreamableHTTP) else MCPServerStreamableHTTP(ts)) for ts in toolsets]
         self.agent = Agent(get_agent_model(), system_prompt=system_prompt, toolsets=toolsets)
+        self.enable_cache = enable_cache
+
+        # Create cached query function if caching is enabled
+        if enable_cache:
+            self._cached_query = task(cache_policy=INPUTS, persist_result=True)(self._run_query)
+
+    def _run_query(self, query: str) -> str:
+        """
+        Internal query execution (no caching)
+        """
+        nest_asyncio.apply()  # NOTE: very important. Because pydantic-ai uses asyncio!
+        result = self.agent.run_sync(query)
+        return result.output
 
     def query(self, query: str) -> str:
         """
+        Run agent query with optional caching
 
         Parameters
         ----------
@@ -38,7 +73,7 @@ class PAIAgent(BaseAgent):
         -------
         str
         """
-
-        nest_asyncio.apply()  # NOTE: very important. Because pydantic-ai uses asyncio!
-        result = self.agent.run_sync(query)
-        return result.output
+        if self.enable_cache:
+            return self._cached_query(query)
+        else:
+            return self._run_query(query)

rdagent/components/agent/context7/__init__.py

@@ -15,7 +15,12 @@ class Agent(PAIAgent):
 
     def __init__(self):
         toolsets = [MCPServerStreamableHTTP(SETTINGS.url, timeout=SETTINGS.timeout)]
-        super().__init__(system_prompt=T(".prompts:system_prompt").r(), toolsets=toolsets)
+
+        super().__init__(
+            system_prompt=T(".prompts:system_prompt").r(),
+            toolsets=toolsets,
+            enable_cache=SETTINGS.enable_cache,
+        )
 
     def _build_enhanced_query(self, error_message: str, full_code: Optional[str] = None) -> str:
         """Build enhanced query using experimental prompt templates."""

rdagent/components/agent/mcp/context7.py

@@ -8,7 +8,7 @@ You can follow the instructions to install it
     cd ~/tmp/context7
     npm install -g bun
     bun i && bun run build
-    bun run dist/index.js --transport http --port 8123 # > bun.out 2>&1 &
+    bun run dist/index.js --transport http --port 8124 # > bun.out 2>&1 &
 """
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -17,8 +17,10 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 class Settings(BaseSettings):
     """Project specific settings."""
 
-    url: str = "http://localhost:8123/mcp"
+    url: str = "http://localhost:8124/mcp"
     timeout: int = 120
+    enable_cache: bool = False
+    # set CONTEXT7_ENABLE_CACHE=true in .env to enable cache
 
     model_config = SettingsConfigDict(
         env_prefix="CONTEXT7_",

requirements.txt

@@ -69,5 +69,5 @@ azureml-mlflow
 types-pytz
 
 # Agent
-pydantic-ai-slim[mcp,openai]
+pydantic-ai-slim[mcp,openai,prefect]
 nest-asyncio

test/oai/test_prefect_cache.py

@@ -0,0 +1,54 @@
+import time
+import unittest
+
+from rdagent.components.agent.context7 import Agent
+
+
+class PydanticTest(unittest.TestCase):
+    """
+    Test Pydantic-AI agent with Prefect caching
+
+    How it works:
+    1. Agent wraps query() with @task(cache_policy=INPUTS) when enable_cache=True
+    2. First call: executes and caches to Prefect server
+    3. Second call with same input: instant cache hit
+    """
+
+    def test_context7_cache(self):
+        """Test that caching works correctly"""
+        query = "pandas read_csv encoding error"
+
+        print("\n" + "=" * 80)
+        print("Testing @task-based caching...")
+        print("=" * 80 + "\n")
+
+        # Create agent once - caching enabled by CONTEXT7_ENABLE_CACHE
+        agent = Agent()
+
+        # First query - will execute and cache
+        print("First query (will execute):")
+        start1 = time.time()
+        res1 = agent.query(query)
+        time1 = time.time() - start1
+
+        print(f"  Time: {time1:.2f}s")
+        print(f"  Length: {len(res1)} chars")
+        print(f"  Preview: {res1[:100]}...\n")
+
+        # Second query - should hit cache (much faster)
+        print("Second query (should hit cache):")
+        start2 = time.time()
+        res2 = agent.query(query)
+        time2 = time.time() - start2
+
+        print(f"  Time: {time2:.2f}s")
+        print(f"  Speedup: {time1/time2:.1f}x faster")
+        print(f"{'='*80}\n")
+
+        self.assertIsNotNone(res1)
+        self.assertGreater(len(res1), 0)
+        self.assertEqual(res1, res2, "Cache must return identical result")
+
+
+if __name__ == "__main__":
+    unittest.main()