From 4f493c8d637dfda42f84af0dc08f8ecfc0501668 Mon Sep 17 00:00:00 2001 From: Jensen Date: Fri, 17 Oct 2025 16:30:59 +0800 Subject: [PATCH] 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 --- .gitignore | 5 ++ rdagent/components/agent/base.py | 47 +++++++++++++--- rdagent/components/agent/context7/__init__.py | 7 ++- rdagent/components/agent/mcp/context7.py | 6 ++- requirements.txt | 2 +- test/oai/test_prefect_cache.py | 54 +++++++++++++++++++ 6 files changed, 111 insertions(+), 10 deletions(-) create mode 100644 test/oai/test_prefect_cache.py diff --git a/.gitignore b/.gitignore index 4d5e0fe1..7f4f8c8c 100644 --- a/.gitignore +++ b/.gitignore @@ -178,3 +178,8 @@ rdagent/app/benchmark/factor/example.json # UI Server resources videos/ static/ + +# AI assistant +.cursor/ +.claude/ +AGENTS.md \ No newline at end of file diff --git a/rdagent/components/agent/base.py b/rdagent/components/agent/base.py index 121a7401..3c15b5a4 100644 --- a/rdagent/components/agent/base.py +++ b/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) diff --git a/rdagent/components/agent/context7/__init__.py b/rdagent/components/agent/context7/__init__.py index 61627329..d8a3794c 100644 --- a/rdagent/components/agent/context7/__init__.py +++ b/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.""" diff --git a/rdagent/components/agent/mcp/context7.py b/rdagent/components/agent/mcp/context7.py index f1fc4ea8..f736944b 100644 --- a/rdagent/components/agent/mcp/context7.py +++ b/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_", diff --git a/requirements.txt b/requirements.txt index adf39826..efd854e2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -69,5 +69,5 @@ azureml-mlflow types-pytz # Agent -pydantic-ai-slim[mcp,openai] +pydantic-ai-slim[mcp,openai,prefect] nest-asyncio diff --git a/test/oai/test_prefect_cache.py b/test/oai/test_prefect_cache.py new file mode 100644 index 00000000..541ab625 --- /dev/null +++ b/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()