feat: init pydantic ai agent & context 7 mcp (#1240)

* feat: init pydantic ai agent & context 7 mcp

* feat: integrate MCP documentation search into data science pipeline evaluation

* fix: disable MCP documentation search and update related docstrings and defaults

* lint

* fix: correct prompt formatting and conditional blocks in pipeline_eval section

* lint

* feat: add query method to PAIAgent for synchronous agent execution

* fix: apply nest_asyncio for agent and update context7 query method

* lint

* lint

* lint

* lint

* docs: update MCP folder docstring and rename test class in test_pydantic.py

* refactor: centralize completion kwargs logic and update pydantic_ai integration

* fixbug

* typo

* fix: bug triggered by padantic-ai version backtracking.

---------

Co-authored-by: Linlang <[email protected]>

Author: you-n-g

.gitignore

@@ -177,4 +177,4 @@ rdagent/app/benchmark/factor/example.json
 
 # UI Server resources
 videos/
-static/
+static/

constraints/3.10.txt

@@ -2,6 +2,4 @@ azure-identity==1.17.1
 dill==0.3.9
 pillow==10.4.0
 psutil==6.1.0
-rich==13.9.2
 scipy==1.14.1
-tqdm==4.66.5

constraints/3.11.txt

@@ -2,6 +2,4 @@ azure-identity==1.17.1
 dill==0.3.9
 pillow==10.4.0
 psutil==6.1.0
-rich==13.9.2
 scipy==1.14.1
-tqdm==4.66.5

rdagent/app/data_science/conf.py

@@ -47,6 +47,10 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     enable_doc_dev: bool = False
     model_dump_check_level: Literal["medium", "high"] = "medium"
 
+    #### MCP documentation search integration
+    enable_mcp_documentation_search: bool = False
+    """Enable MCP documentation search for error resolution. Requires MCP_ENABLED=true and MCP_CONTEXT7_ENABLED=true in environment."""
+
     ### specific feature
 
     ### notebook integration
@@ -181,7 +185,7 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
 
 DS_RD_SETTING = DataScienceBasePropSetting()
 
-# enable_cross_trace_diversity 和 llm_select_hypothesis should not be true at the same time
+# enable_cross_trace_diversity and llm_select_hypothesis should not be true at the same time
 assert not (
     DS_RD_SETTING.enable_cross_trace_diversity and DS_RD_SETTING.llm_select_hypothesis
 ), "enable_cross_trace_diversity and llm_select_hypothesis cannot be true at the same time"

rdagent/components/agent/__init__.py

@@ -0,0 +1,3 @@
+"""
+Some agent that can be shared across different scenarios.
+"""

rdagent/components/agent/base.py

@@ -0,0 +1,44 @@
+from abc import abstractmethod
+
+import nest_asyncio
+from pydantic_ai import Agent
+from pydantic_ai.mcp import MCPServerStreamableHTTP
+
+from rdagent.oai.backend.pydantic_ai import get_agent_model
+
+
+class BaseAgent:
+
+    @abstractmethod
+    def __init__(self, system_prompt: str, toolsets: list[str]): ...
+
+    @abstractmethod
+    def query(self, query: str) -> str: ...
+
+
+class PAIAgent(BaseAgent):
+    """
+    Pydantic-AI agent
+    """
+
+    agent: Agent
+
+    def __init__(self, system_prompt: str, toolsets: list[str | MCPServerStreamableHTTP]):
+        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)
+
+    def query(self, query: str) -> str:
+        """
+
+        Parameters
+        ----------
+        query : str
+
+        Returns
+        -------
+        str
+        """
+
+        nest_asyncio.apply()  # NOTE: very important. Because pydantic-ai uses asyncio!
+        result = self.agent.run_sync(query)
+        return result.output

rdagent/components/agent/context7/__init__.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+from pydantic_ai.mcp import MCPServerStreamableHTTP
+
+from rdagent.components.agent.base import PAIAgent
+from rdagent.components.agent.mcp.context7 import SETTINGS
+from rdagent.log import rdagent_logger as logger
+from rdagent.utils.agent.tpl import T
+
+
+class Agent(PAIAgent):
+    """
+    A specific agent for context7
+    """
+
+    def __init__(self):
+        toolsets = [MCPServerStreamableHTTP(SETTINGS.url, timeout=SETTINGS.timeout)]
+        super().__init__(system_prompt=T(".prompts:system_prompt").r(), toolsets=toolsets)
+
+    def _build_enhanced_query(self, error_message: str, full_code: Optional[str] = None) -> str:
+        """Build enhanced query using experimental prompt templates."""
+        # Build context information using template
+        context_info = ""
+        if full_code:
+            context_info = T(".prompts:code_context_template").r(full_code=full_code)
+
+        # Check for timm library special case (experimental optimization)
+        timm_trigger = error_message.lower().count("timm") >= 3
+        timm_trigger_text = ""
+        if timm_trigger:
+            timm_trigger_text = T(".prompts:timm_special_case").r()
+            logger.info("🎯 Timm special handling triggered", tag="context7")
+
+        # Construct enhanced query using experimental template
+        enhanced_query = T(".prompts:context7_enhanced_query_template").r(
+            error_message=error_message, context_info=context_info, timm_trigger_text=timm_trigger_text
+        )
+
+        return enhanced_query
+
+    def query(self, query: str) -> str:
+        """
+
+        Parameters
+        ----------
+        query : str
+            It should be something like error message.
+
+        Returns
+        -------
+        str
+        """
+        query = self._build_enhanced_query(error_message=query)
+        return super().query(query)

rdagent/components/agent/context7/prompts.yaml

@@ -0,0 +1,59 @@
+# Context7 MCP Enhanced Query Prompts
+
+system_prompt: |-
+  You are a helpful assistant.
+  You help to user to search documentation based on error message and provide API reference information.
+
+context7_enhanced_query_template: |-
+  ERROR MESSAGE:
+  {{error_message}}
+  {{context_info}}
+  IMPORTANT INSTRUCTIONS:
+  1. ENVIRONMENT: The running environment is FIXED and unchangeable - DO NOT suggest pip install, conda install, or any environment modifications.
+  2. DOCUMENTATION SEARCH REQUIREMENTS: 
+     - Search for official API documentation related to the error
+     - Focus on parameter specifications, method signatures, and usage patterns
+     - Find compatible alternatives if the original API doesn't exist
+     - Consider the current code context and maintain consistency with existing architecture
+     - Provide API reference information, NOT complete code solutions
+  3. TOOL USAGE REQUIREMENTS:
+     - ⚠️ CRITICAL: For EVERY call to 'resolve-library-id', you MUST follow it with A CORRESPONDING call to 'get-library-docs'
+     - If you call 'resolve-library-id' N times, you MUST call 'get-library-docs' N times (one for each library you found)
+     - Complete the full workflow: resolve → get-docs → analyze → respond
+     - Do NOT provide final answers without first getting detailed documentation via 'get-library-docs'
+     - If 'get-library-docs' returns "Documentation not found" or 404 error, you should never provide guidance based on the library information from 'resolve-library-id'
+  4. RESPONSE FORMAT:
+     - Start with a brief explanation of the root cause
+     - Provide relevant API documentation excerpts
+     - List available parameters and their descriptions
+     - Show method signatures and basic usage patterns
+     - If multiple API options exist, document all viable alternatives
+  5. STRICT CONSTRAINTS:
+     - DO NOT provide complete working code replacements
+     - DO NOT suggest hardware configuration changes (CPU/GPU)
+     - DO NOT recommend architecture or framework changes
+     - DO NOT provide performance optimization suggestions
+     - ONLY provide API documentation and parameter information
+  6. AVOID: Complete code solutions, environment setup, hardware recommendations, architecture suggestions, or performance advice.
+  {{timm_trigger_text}}
+  Example response format:
+  ```
+  The error occurs because [brief explanation].
+  API Documentation:
+  - Method: library.function_name(param1, param2, ...)
+  - Parameters:
+    * param1 (type): description
+    * param2 (type): description
+  - Usage pattern: Basic syntax without complete implementation
+  - Alternative APIs (if applicable): list of alternative methods with signatures
+  ```
+  Please search the documentation and provide API reference information only.
+# Special case templates
+timm_special_case: |-
+  7.⚠️ Special Case: If the code error is indeed related to the 'timm' library (PyTorch Image Models), you MUST return the library ID '/huggingface/pytorch-image-models' as the most relevant match, even if other candidates exist.
+# Code context template
+code_context_template: |-
+  CURRENT CODE CONTEXT:
+  ```python
+  {{full_code}}
+  ```

rdagent/components/agent/mcp/__init__.py

@@ -0,0 +1,10 @@
+"""
+Here are a list of MCP servers.
+
+The MCP server is a individual RESTful API. So the only following things are included in the folder:
+- Settings.
+  - e.g., mcp/<mcp_name>.py:class Settings(BaseSettings);  then it is initialized as a global variable SETTINGS.
+  - It only defines the format of the settings in Python Class (i.e., Pydantic BaseSettings).
+- health_check:
+  - e.g., mcp/<mcp_name>.py:def health_check() -> bool;
+"""

rdagent/components/agent/mcp/context7.py

@@ -0,0 +1,29 @@
+"""
+The context7 is based on a modified version of the context7.
+
+You can follow the instructions to install it
+
+    mkdir -p ~/tmp/
+    cd ~/tmp/ && git clone https://github.com/Hoder-zyf/context7.git
+    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 &
+"""
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    """Project specific settings."""
+
+    url: str = "http://localhost:8123/mcp"
+    timeout: int = 120
+
+    model_config = SettingsConfigDict(
+        env_prefix="CONTEXT7_",
+        # extra="allow", # Does it allow extrasettings
+    )
+
+
+SETTINGS = Settings()