diff --git a/rdagent/app/data_science/conf.py b/rdagent/app/data_science/conf.py index 14f107d5..2d8ec626 100644 --- a/rdagent/app/data_science/conf.py +++ b/rdagent/app/data_science/conf.py @@ -175,6 +175,9 @@ class DataScienceBasePropSetting(KaggleBasePropSetting): enable_generate_unique_hypothesis: bool = False """Enable generate unique hypothesis. If True, generate unique hypothesis for each component. If False, generate unique hypothesis for each component.""" + enable_research_rag: bool = False + """Enable research RAG for hypothesis generation.""" + #### hypothesis critique and rewrite enable_hypo_critique_rewrite: bool = False """Enable hypothesis critique and rewrite stages for improving hypothesis quality""" diff --git a/rdagent/components/agent/context7/__init__.py b/rdagent/components/agent/context7/__init__.py index d8a3794c..b7ae2a08 100644 --- a/rdagent/components/agent/context7/__init__.py +++ b/rdagent/components/agent/context7/__init__.py @@ -3,7 +3,7 @@ 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.components.agent.context7.conf import SETTINGS from rdagent.log import rdagent_logger as logger from rdagent.utils.agent.tpl import T diff --git a/rdagent/components/agent/mcp/context7.py b/rdagent/components/agent/context7/conf.py similarity index 100% rename from rdagent/components/agent/mcp/context7.py rename to rdagent/components/agent/context7/conf.py diff --git a/rdagent/components/agent/rag/__init__.py b/rdagent/components/agent/rag/__init__.py new file mode 100644 index 00000000..b3d7c364 --- /dev/null +++ b/rdagent/components/agent/rag/__init__.py @@ -0,0 +1,17 @@ +from pydantic_ai.mcp import MCPServerStreamableHTTP + +from rdagent.components.agent.base import PAIAgent +from rdagent.components.agent.rag.conf import SETTINGS +from rdagent.utils.agent.tpl import T + + +class Agent(PAIAgent): + """ + A specific agent for RAG + """ + + def __init__(self, system_prompt: str | None = None): + toolsets = [MCPServerStreamableHTTP(SETTINGS.url, timeout=SETTINGS.timeout)] + if system_prompt is None: + system_prompt = "You are a Retrieval-Augmented Generation (RAG) agent. Use the retrieved documents to answer the user's queries accurately and concisely." + super().__init__(system_prompt=system_prompt, toolsets=toolsets) diff --git a/rdagent/components/agent/rag/conf.py b/rdagent/components/agent/rag/conf.py new file mode 100644 index 00000000..e2d8d59a --- /dev/null +++ b/rdagent/components/agent/rag/conf.py @@ -0,0 +1,22 @@ +""" +Settings for RAG agent. + +TODO: how run the RAG mcp server +""" + +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + """Project specific settings.""" + + url: str = "http://localhost:8124/mcp" + timeout: int = 120 + + model_config = SettingsConfigDict( + env_prefix="RAG_", + # extra="allow", # Does it allow extrasettings + ) + + +SETTINGS = Settings() diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml index d7a5567b..144485d6 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml +++ b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml @@ -294,6 +294,11 @@ hypothesis_gen: # Identified Challenges{% if enable_idea_pool %} with Sampled Ideas{% endif %} {{ problems }} + {% if knowledge %} + # Some reference knowledge from the community + {{ knowledge }} + {% endif %} + hypothesis_critique: system: |- {% include "scenarios.data_science.share:scen.role" %} diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py index 0cac34e6..8a0343f2 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py @@ -9,6 +9,7 @@ import pandas as pd from pydantic import BaseModel, Field from rdagent.app.data_science.conf import DS_RD_SETTING +from rdagent.components.agent.rag import Agent as RAGAgent from rdagent.components.coder.data_science.ensemble.exp import EnsembleTask from rdagent.components.coder.data_science.feature.exp import FeatureTask from rdagent.components.coder.data_science.model.exp import ModelTask @@ -645,12 +646,24 @@ class DSProposalV2ExpGen(ExpGen): sibling_hypotheses=sibling_hypotheses, former_user_instructions_str=str(former_user_instructions) if former_user_instructions else None, ) + + # knowledge retrieval + if DS_RD_SETTING.enable_research_rag: + rag_agent = RAGAgent( + system_prompt="""You are a helpful assistant. +You help users retrieve relevant knowledge from community discussions and public code.""" + ) + knowledge = rag_agent.query(problem_formatted_str) + else: + knowledge = None + user_prompt = T(".prompts_v2:hypothesis_gen.user").r( scenario_desc=scenario_desc, exp_and_feedback_list_desc=exp_feedback_list_desc, sota_exp_desc=sota_exp_desc, problems=problem_formatted_str, enable_idea_pool=enable_idea_pool, + knowledge=knowledge, ) response = APIBackend().build_messages_and_create_chat_completion( user_prompt=user_prompt,