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 <Lv.Linlang@hotmail.com>

.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()

rdagent/components/coder/CoSTEER/evaluators.py

@@ -76,6 +76,27 @@ class CoSTEERSingleFeedback(Feedback):
                 raise ValueError(f"'{attr}' must be a string, not {type(data[attr])}")
         return data
 
+    @classmethod
+    def merge(cls, feedback_li: list["CoSTEERSingleFeedback"]) -> "CoSTEERSingleFeedback":
+        # NOTE:
+        # Here we don't know the detailed design of each feedback, we just know they are CoSTEERSingleFeedback
+        # So we merge them only based on CoSTEERSingleFeedback's attributes
+        # **So some information may be lost when we have different types of feedbacks**
+        # If you have more sophisticated sub class of CoSTEERSingleFeedback, you should override this method
+        # to avoid the loss of information.
+
+        fb = deepcopy(feedback_li[0])
+
+        # for all the evaluators, aggregate the final_decision from `task_id`
+        fb.final_decision = all(fb.final_decision for fb in feedback_li)
+        for attr in "execution", "return_checking", "code":
+            setattr(
+                fb,
+                attr,
+                "\n\n".join([getattr(_fb, attr) for _fb in feedback_li if getattr(_fb, attr) is not None]),
+            )
+        return fb
+
     def __str__(self) -> str:
         return f"""------------------Execution------------------
 {self.execution}
@@ -230,7 +251,18 @@ class CoSTEERMultiEvaluator(CoSTEEREvaluator):
         **kwargs,
     ) -> CoSTEERMultiFeedback:
         eval_l = self.single_evaluator if isinstance(self.single_evaluator, list) else [self.single_evaluator]
+
+        # 1) Evaluate each sub_task
         task_li_feedback_li = []
+        # task_li_feedback_li: List[List[CoSTEERSingleFeedback]]
+        # Example:
+        # If there are 2 evaluators and 3 sub_tasks in evo, and each evaluator's evaluate returns a list of 3 CoSTEERSingleFeedbacks,
+        # Then task_li_feedback_li will be:
+        # [
+        #   [feedback_1_1, feedback_1_2, feedback_1_3],  # results from the 1st evaluator for all sub_tasks
+        #   [feedback_2_1, feedback_2_2, feedback_2_3],  # results from the 2nd evaluator for all sub_tasks
+        # ]
+        # Where feedback_i_j is the feedback from the i-th evaluator for the j-th sub_task.
         for ev in eval_l:
             multi_implementation_feedback = multiprocessing_wrapper(
                 [
@@ -248,27 +280,22 @@ class CoSTEERMultiEvaluator(CoSTEEREvaluator):
                 n=RD_AGENT_SETTINGS.multi_proc_n,
             )
             task_li_feedback_li.append(multi_implementation_feedback)
-        # merge the feedbacks
-        merged_task_feedback = []
-        for task_id, fb in enumerate(task_li_feedback_li[0]):
-            fb = deepcopy(fb)  # deep copy to make it more robust
 
-            fb.final_decision = all(
+        # 2) merge the feedbacks along the sub_tasks to aggregate the multiple evaluation feedbacks
-                task_li_feedback[task_id].final_decision for task_li_feedback in task_li_feedback_li
+        merged_task_feedback = []
-            )
+        # task_li_feedback_li[0] is a list of feedbacks of different tasks for the 1st evaluator
-            for attr in "execution", "return_checking", "code":
+        for task_id, fb in enumerate(task_li_feedback_li[0]):
-                setattr(
+            fb = fb.merge([fb_li[task_id] for fb_li in task_li_feedback_li])
-                    fb,
-                    attr,
-                    "\n\n".join(
-                        [
-                            getattr(task_li_feedback[task_id], attr)
-                            for task_li_feedback in task_li_feedback_li
-                            if getattr(task_li_feedback[task_id], attr) is not None
-                        ]
-                    ),
-                )
             merged_task_feedback.append(fb)
+        # merged_task_feedback: List[CoSTEERSingleFeedback]
+        # Example:
+        # [
+        #   CoSTEERSingleFeedback(final_decision=True, execution="...", return_checking="...", code="..."),
+        #   CoSTEERSingleFeedback(final_decision=False, execution="...", return_checking="...", code="..."),
+        #   ...
+        # ]
+        # Each element corresponds to the merged feedback for one sub-task across all evaluators.
+        # merged_task_feedback[i] is the merged feedback for the i-th sub_task
 
         final_decision = [
             None if single_feedback is None else single_feedback.final_decision

rdagent/components/coder/data_science/pipeline/eval.py

@@ -2,11 +2,13 @@
 # (GPT) if it aligns with the spec & rationality of the spec.
 import json
 import re
+from dataclasses import dataclass
 from pathlib import Path
 
 import pandas as pd
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.components.agent.context7 import Agent as DocAgent
 from rdagent.components.coder.CoSTEER import CoSTEERMultiFeedback
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEEREvaluator,
@@ -19,13 +21,102 @@ from rdagent.components.coder.data_science.conf import get_clear_ws_cmd, get_ds_
 from rdagent.components.coder.data_science.share.notebook import NotebookConverter
 from rdagent.components.coder.data_science.utils import remove_eda_part
 from rdagent.core.experiment import FBWorkspace, Task
+from rdagent.log import rdagent_logger as logger
 from rdagent.scenarios.data_science.test_eval import get_test_eval
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.agent.workflow import build_cls_from_json_with_retry
 
 DIRNAME = Path(__file__).absolute().resolve().parent
 
-PipelineSingleFeedback = CoSTEERSingleFeedback
+
+@dataclass
+class DSCoderFeedback(CoSTEERSingleFeedback):
+    """
+    Feedback for Data Science CoSTEER evaluation.
+    This feedback is used to evaluate the code and execution of the Data Science CoSTEER task.
+    """
+
+    requires_documentation_search: bool = False
+    error_message: str | None = None
+
+    @staticmethod
+    def val_and_update_init_dict(data: dict) -> dict:
+        # First call parent class validation method to handle base fields
+        data = CoSTEERSingleFeedback.val_and_update_init_dict(data)
+
+        # Validate new fields
+        if "requires_documentation_search" in data:
+            if isinstance(data["requires_documentation_search"], str):
+                if data["requires_documentation_search"] == "false" or data["requires_documentation_search"] == "False":
+                    data["requires_documentation_search"] = False
+                elif data["requires_documentation_search"] == "true" or data["requires_documentation_search"] == "True":
+                    data["requires_documentation_search"] = True
+                else:
+                    raise ValueError(
+                        f"'requires_documentation_search' string value must be 'true', 'True', 'false', or 'False', not '{data['requires_documentation_search']}'"
+                    )
+            elif data["requires_documentation_search"] is not None and not isinstance(
+                data["requires_documentation_search"], bool
+            ):
+                raise ValueError(
+                    f"'requires_documentation_search' must be a boolean, string, or None, not {type(data['requires_documentation_search'])}"
+                )
+
+        if "error_message" in data:
+            if data["error_message"] is not None and not isinstance(data["error_message"], str):
+                raise ValueError(f"'error_message' must be a string or None, not {type(data['error_message'])}")
+
+        return data
+
+    def __str__(self) -> str:
+        base_str = super().__str__()
+
+        if self.requires_documentation_search is not None:
+            base_str += f"-------------------Documentation Search Required------------------\n{self.requires_documentation_search}\n"
+
+        if self.error_message is not None:
+            # Check if error_message contains Context7 documentation results
+            if "### API Documentation Reference:" in self.error_message:
+                base_str += f"-------------------Error Analysis & Documentation Search Results ------------------\n{self.error_message}\n"
+            else:
+                base_str += f"-------------------Error Message------------------\n{self.error_message}\n"
+
+        return base_str
+
+    @classmethod
+    def merge(cls, feedback_li: list[CoSTEERSingleFeedback]) -> "DSCoderFeedback":
+        # Call parent class merge method to handle base fields
+        merged_fb = super().merge(feedback_li)
+
+        # Convert to DSCoderFeedback type if needed
+        if not isinstance(merged_fb, DSCoderFeedback):
+            merged_fb = DSCoderFeedback(
+                execution=merged_fb.execution,
+                return_checking=merged_fb.return_checking,
+                code=merged_fb.code,
+                final_decision=merged_fb.final_decision,
+            )
+
+        # Merge error_message fields
+        error_messages = [
+            fb.error_message for fb in feedback_li if isinstance(fb, DSCoderFeedback) and fb.error_message is not None
+        ]
+        if error_messages:
+            merged_fb.error_message = "\n\n".join(error_messages)
+
+        # Merge requires_documentation_search fields (True if any is True)
+        requires_search = [
+            fb.requires_documentation_search
+            for fb in feedback_li
+            if isinstance(fb, DSCoderFeedback) and fb.requires_documentation_search is not None
+        ]
+        if requires_search:
+            merged_fb.requires_documentation_search = any(requires_search)
+
+        return merged_fb
+
+
+PipelineSingleFeedback = DSCoderFeedback  # Only for compatible
 PipelineMultiFeedback = CoSTEERMultiFeedback
 
 
@@ -51,6 +142,8 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
                 execution="This task has failed too many times, skip implementation.",
                 return_checking="This task has failed too many times, skip implementation.",
                 code="This task has failed too many times, skip implementation.",
+                error_message="This task has failed too many times, skip implementation.",
+                requires_documentation_search=False,
                 final_decision=False,
             )
 
@@ -177,6 +270,9 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
         else:
             eda_output = implementation.file_dict.get("EDA.md", None)
 
+        # extract enable_mcp_documentation_search from data science configuration
+        enable_mcp_documentation_search = DS_RD_SETTING.enable_mcp_documentation_search
+
         queried_similar_successful_knowledge = (
             queried_knowledge.task_to_similar_task_successful_knowledge[target_task.get_task_information()]
             if queried_knowledge is not None
@@ -186,6 +282,7 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
         system_prompt = T(".prompts:pipeline_eval.system").r(
             is_sub_enabled=test_eval.is_sub_enabled(self.scen.competition),
             debug_mode=DS_RD_SETTING.sample_data_by_LLM,
+            enable_mcp_documentation_search=enable_mcp_documentation_search,
             mle_check=DS_RD_SETTING.sample_data_by_LLM,
             queried_similar_successful_knowledge=queried_similar_successful_knowledge,
         )
@@ -205,6 +302,35 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
             user_prompt=user_prompt,
             init_kwargs_update_func=PipelineSingleFeedback.val_and_update_init_dict,
         )
+
+        # judge whether we should perform documentation search
+        do_documentation_search = enable_mcp_documentation_search and wfb.requires_documentation_search
+
+        if do_documentation_search:
+            # Use MCPAgent for clean, user-friendly interface
+            try:
+                # Create agent targeting Context7 service - model config comes from mcp_config.json
+                doc_agent = DocAgent()
+
+                # Synchronous query - perfect for evaluation context
+                if wfb.error_message:  # Type safety check
+                    context7_result = doc_agent.query(query=wfb.error_message)
+
+                    if context7_result:
+                        logger.info("Context7: Documentation search completed successfully")
+                        wfb.error_message += f"\n\n### API Documentation Reference:\nThe following API documentation was retrieved based on the error. This provides factual information about API changes or parameter specifications only:\n\n{context7_result}"
+                    else:
+                        logger.warning("Context7: Documentation search failed or no results found")
+                else:
+                    logger.warning("Context7: No error message to search for")
+
+            # TODO: confirm what exception will be raised when timeout
+            # except concurrent.futures.TimeoutError:
+            #     logger.error("Context7: Query timed out after 180 seconds")
+            except Exception as e:
+                error_msg = str(e) if str(e) else type(e).__name__
+                logger.error(f"Context7: Query failed - {error_msg}")
+
         if score_ret_code != 0 and wfb.final_decision is True:
             wfb.final_decision = False
             wfb.return_checking += "\n" + score_check_text

rdagent/components/coder/data_science/pipeline/prompts.yaml

@@ -235,7 +235,11 @@ pipeline_eval:
     - If the code execute successfully:
       - Proceed to Step 2.
     - If the code does not execute successfully:
-      - Set the "final_decision" to false and write complete analysis in the "execution" field.
+      - Set the "final_decision" to false.
+      {% if enable_mcp_documentation_search %}
+      - Given that my package/environment is fixed and unchangeable, first you should go through the code and the execution output,if the problem could be solved by looking up the official documentation to confirm feature/API availability, compatible usage, or official alternatives in the fixed environment, set the "requires_documentation_search" to true.
+      {% endif %}
+      - Write complete analysis in the "execution" field.
 
     ### Competition Alignment
     - Goal: Confirm strict adherence to the competition's evaluation rules and experimental setup.
@@ -309,10 +313,14 @@ pipeline_eval:
     Please respond with your feedback in the following JSON format without anything else.
     ```json
     {
-        "execution": "Describe whether the code executed successfully. Include any errors or issues encountered, and append all error messages and full traceback details without summarizing or omitting any information. If errors occurred, analyze the root causes: (1) Are they fundamental algorithmic/approach issues, or (2) Implementation details that can be easily fixed, or (3) Environment/dependency problems?",
+        {% if enable_mcp_documentation_search %}
+        "requires_documentation_search": <true/false>,
+        {% endif %}"execution": "Describe whether the code executed successfully. Include any errors or issues encountered, and append all error messages and full traceback details without summarizing or omitting any information. If errors occurred, analyze the root causes: (1) Are they fundamental algorithmic/approach issues, or (2) Implementation details that can be easily fixed, or (3) Environment/dependency problems?",
         "return_checking": "Examine the generated files by cross-referencing the code logic and stdout output. Verify: (1) Format matches required submission format (index, column names, CSV content); (2) **File generation authenticity**: Is the file genuinely produced by successful model execution, or is it a result of exception handling/fallback mechanisms? Cite specific code sections and stdout evidence.",
         "code": "Begin explicitly with [Code analysis] or [Evaluation error]. Provide structured analysis: (1) **Technical Appropriateness**: Does the chosen approach (algorithms, data processing, validation strategy) match this problem's data characteristics and competition requirements? (2) **Effective Components**: What specific parts work well and why are they effective for this problem type? (3) **Issues & Improvements**: Identify concrete problems and suggest actionable improvement directions (without providing actual code). (4) **Code Quality**: Assess readability, structure, and adherence to specifications.",
-        "final_decision": <true/false>
+        {% if enable_mcp_documentation_search %}
+        "error_message": "If the code execution has problems, extract the error information in the following format, otherwise set to empty string: ### TRACEBACK: <full relevant traceback extracted from execution output> ### SUPPLEMENTARY_INFO: <only if TRACEBACK is unclear - copy exact code fragments: import statements, variable=value assignments, function calls with parameters as they appear in code>",
+        {% endif %}"final_decision": <true/false>
     }
     ```
 

rdagent/oai/backend/__init__.py

@@ -1,2 +1,2 @@
-from .deprec import DeprecBackend
+from .deprec import DeprecBackend  # type: ignore[attr-defined]
 from .litellm import LiteLLMAPIBackend

rdagent/oai/backend/deprec.py

@@ -1,3 +1,4 @@
+# type: ignore
 from __future__ import annotations
 
 import inspect

rdagent/oai/backend/litellm.py

@@ -1,5 +1,5 @@
 import copyreg
-from typing import Any, Literal, Optional, Type, Union, cast
+from typing import Any, Literal, Optional, Type, TypedDict, Union, cast
 
 import numpy as np
 from litellm import (
@@ -85,6 +85,44 @@ class LiteLLMAPIBackend(APIBackend):
         response_list = [data["embedding"] for data in response.data]
         return response_list
 
+    class CompleteKwargs(TypedDict):
+        model: str
+        temperature: float
+        max_tokens: int | None
+        reasoning_effort: Literal["low", "medium", "high"] | None
+
+    def get_complete_kwargs(self) -> CompleteKwargs:
+        """
+        return several key settings for completion
+        getting these values from settings makes it easier to adapt to backend calls in agent systems.
+        """
+        # Call LiteLLM completion
+        model = LITELLM_SETTINGS.chat_model
+        temperature = LITELLM_SETTINGS.chat_temperature
+        max_tokens = LITELLM_SETTINGS.chat_max_tokens
+        reasoning_effort = LITELLM_SETTINGS.reasoning_effort
+
+        if LITELLM_SETTINGS.chat_model_map:
+            for t, mc in LITELLM_SETTINGS.chat_model_map.items():
+                if t in logger._tag:
+                    model = mc["model"]
+                    if "temperature" in mc:
+                        temperature = float(mc["temperature"])
+                    if "max_tokens" in mc:
+                        max_tokens = int(mc["max_tokens"])
+                    if "reasoning_effort" in mc:
+                        if mc["reasoning_effort"] in ["low", "medium", "high"]:
+                            reasoning_effort = cast(Literal["low", "medium", "high"], mc["reasoning_effort"])
+                        else:
+                            reasoning_effort = None
+                    break
+        return self.CompleteKwargs(
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            reasoning_effort=reasoning_effort,
+        )
+
     def _create_chat_completion_inner_function(  # type: ignore[no-untyped-def] # noqa: C901, PLR0912, PLR0915
         self,
         messages: list[dict[str, Any]],
@@ -109,34 +147,15 @@ class LiteLLMAPIBackend(APIBackend):
 
         if LITELLM_SETTINGS.log_llm_chat_content:
             logger.info(self._build_log_messages(messages), tag="llm_messages")
-        # Call LiteLLM completion
-        model = LITELLM_SETTINGS.chat_model
-        temperature = LITELLM_SETTINGS.chat_temperature
-        max_tokens = LITELLM_SETTINGS.chat_max_tokens
-        reasoning_effort = LITELLM_SETTINGS.reasoning_effort
 
-        if LITELLM_SETTINGS.chat_model_map:
+        complete_kwargs = self.get_complete_kwargs()
-            for t, mc in LITELLM_SETTINGS.chat_model_map.items():
+        model = complete_kwargs["model"]
-                if t in logger._tag:
+
-                    model = mc["model"]
-                    if "temperature" in mc:
-                        temperature = float(mc["temperature"])
-                    if "max_tokens" in mc:
-                        max_tokens = int(mc["max_tokens"])
-                    if "reasoning_effort" in mc:
-                        if mc["reasoning_effort"] in ["low", "medium", "high"]:
-                            reasoning_effort = cast(Literal["low", "medium", "high"], mc["reasoning_effort"])
-                        else:
-                            reasoning_effort = None
-                    break
         response = completion(
-            model=model,
             messages=messages,
             stream=LITELLM_SETTINGS.chat_stream,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            reasoning_effort=reasoning_effort,
             max_retries=0,
+            **complete_kwargs,
             **kwargs,
         )
         logger.info(f"{LogColors.GREEN}Using chat model{LogColors.END} {model}", tag="llm_messages")

rdagent/oai/backend/pydantic_ai.py

@@ -0,0 +1,63 @@
+"""
+Adapter tools for pydantic-ai
+"""
+
+import os
+
+from litellm.utils import get_llm_provider
+from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
+from pydantic_ai.providers.litellm import LiteLLMProvider
+
+from rdagent.oai.backend.litellm import LiteLLMAPIBackend
+from rdagent.oai.llm_conf import LLM_SETTINGS
+from rdagent.oai.llm_utils import APIBackend
+
+# NOTE:
+# LiteLLM's code is not well orgnized.
+# we can't reuse any component to map the provider to the env name
+# So we have to hardcode on here.
+PROVIDER_TO_ENV_MAP = {
+    "openai": "OPENAI",
+    "azure_ai": "AZURE_AI",
+    "azure": "AZURE",
+    "litellm_proxy": "LITELLM_PROXY",
+}
+
+
+def get_agent_model() -> OpenAIChatModel:
+    """
+    Converting LiteLLM to a pydantic-ai model. So you can use like this
+
+    .. code-block:: python
+
+        from rdagent.oai.backend.pydantic_ai import get_agent_model
+        model = get_agent_model()
+        agent = Agent(model)
+
+    """
+    backend = APIBackend()
+    assert isinstance(backend, LiteLLMAPIBackend), "Only LiteLLMAPIBackend is supported"
+
+    compl_kwargs = backend.get_complete_kwargs()
+
+    selected_model = compl_kwargs["model"]
+
+    _, custom_llm_provider, _, _ = get_llm_provider(selected_model)
+    assert (
+        custom_llm_provider in PROVIDER_TO_ENV_MAP
+    ), f"Provider {custom_llm_provider} not supported. Please add it into `PROVIDER_TO_ENV_MAP`"
+    prefix = PROVIDER_TO_ENV_MAP[custom_llm_provider]
+    api_key = os.getenv(f"{prefix}_API_KEY", None)
+    api_base = os.getenv(f"{prefix}_API_BASE", None)
+
+    kwargs = {
+        "openai_reasoning_effort": compl_kwargs.get("reasoning_effort"),
+        "max_tokens": compl_kwargs.get("max_tokens"),
+        "temperature": compl_kwargs.get("temperature"),
+    }
+    if compl_kwargs.get("max_tokens") is None:
+        kwargs["max_tokens"] = LLM_SETTINGS.chat_max_tokens
+    settings = OpenAIChatModelSettings(**kwargs)
+    return OpenAIChatModel(
+        selected_model, provider=LiteLLMProvider(api_base=api_base, api_key=api_key), settings=settings
+    )

requirements.txt

@@ -67,3 +67,7 @@ genson
 mlflow
 azureml-mlflow
 types-pytz
+
+# Agent
+pydantic-ai-slim[mcp,openai]
+nest-asyncio

test/oai/test_pydantic.py

@@ -0,0 +1,15 @@
+import unittest
+
+from rdagent.components.agent.context7 import Agent
+
+
+class PydanticTest(unittest.TestCase):
+
+    def test_context7(self):
+        context7a = Agent()
+        res = context7a.query("pandas read_csv encoding error")
+        print(res)
+
+
+if __name__ == "__main__":
+    unittest.main()