fix: fix some bugs in rag (#399)

* fix some bugs in rag

* fix a bug in costeer

rdagent/app/kaggle/loop.py

@@ -19,14 +19,14 @@ from rdagent.core.scenario import Scenario
 from rdagent.core.utils import import_class
 from rdagent.log import rdagent_logger as logger
 from rdagent.log.time import measure_time
-from rdagent.scenarios.kaggle.experiment.utils import python_files_to_notebook
-from rdagent.scenarios.kaggle.kaggle_crawler import download_data
-from rdagent.scenarios.kaggle.proposal.proposal import (
+from rdagent.scenarios.kaggle.experiment.scenario import (
     KG_ACTION_FEATURE_ENGINEERING,
     KG_ACTION_FEATURE_PROCESSING,
     KG_ACTION_MODEL_FEATURE_SELECTION,
-    KGTrace,
 )
+from rdagent.scenarios.kaggle.experiment.utils import python_files_to_notebook
+from rdagent.scenarios.kaggle.kaggle_crawler import download_data
+from rdagent.scenarios.kaggle.proposal.proposal import KGTrace
 
 
 class KaggleRDLoop(RDLoop):

rdagent/components/coder/factor_coder/CoSTEER/evolving_strategy.py

@@ -90,6 +90,7 @@ class MultiProcessEvolvingStrategy(EvolvingStrategy):
             ],
             n=RD_AGENT_SETTINGS.multi_proc_n,
         )
+        from rdagent.components.coder.factor_coder.factor import FactorFBWorkspace
 
         for index, target_index in enumerate(to_be_finished_task_index):
             if evo.sub_workspace_list[target_index] is None:

rdagent/components/proposal/model_proposal.py

@@ -39,7 +39,7 @@ class ModelHypothesisGen(HypothesisGen):
             Environment(undefined=StrictUndefined)
             .from_string(ModelHypothesisGen.prompts["hypothesis_gen"]["system_prompt"])
             .render(
-                targets="feature engineering and model building",
+                targets="model tuning",
                 scenario=self.scen.get_scenario_all_desc(),
                 hypothesis_output_format=context_dict["hypothesis_output_format"],
                 hypothesis_specification=context_dict["hypothesis_specification"],
@@ -49,7 +49,7 @@ class ModelHypothesisGen(HypothesisGen):
             Environment(undefined=StrictUndefined)
             .from_string(ModelHypothesisGen.prompts["hypothesis_gen"]["user_prompt"])
             .render(
-                targets="feature engineering and model building",
+                targets="model tuning",
                 RAG=context_dict["RAG"],
             )
         )

rdagent/components/proposal/prompts.yaml

@@ -1,19 +1,24 @@
 hypothesis_gen:
   system_prompt: |-
-    The user is trying to generate new hypothesis on the {{targets}} in data-driven research and development.
-    The {{targets}} are used in a certain scenario, the scenario is as follows:
-    {{ scenario }}
-    The user has made several hypothesis on this scenario and did several evaluation on them. The user will provide this information to you. Check if a new hypothesis has already been proposed. If it is already generated and you agree with it, just use it. If you don't agree, generate a better one.
+    The user is working on generating new hypotheses for the {{targets}} in a data-driven research and development process. 
+    The {{targets}} are used in the following scenario:
+    {{scenario}}
+    The user has already proposed several hypotheses and conducted evaluations on them. This information will be provided to you. Your task is to check whether a similar hypothesis has already been generated. 
+    If one exists and you agree with it, feel free to use it. If you disagree, please generate an improved version.
     {% if hypothesis_specification %}
-    To help you generate new hypothesis, the user has prepared some additional information for you. You should use this information to help generate new {{targets}}.
-    Here are the specifications: {{ hypothesis_specification }}
+    To assist you in formulating new hypotheses, the user has provided some additional information: {{hypothesis_specification}}.
+    **Important:** If the hypothesis_specification outlines the next steps you need to follow, ensure you adhere to those instructions.
     {% endif %}
-    Please generate the output following the format and specifications below:
+    Please generate the output using the following format and specifications:
     {{ hypothesis_output_format }}
 
   user_prompt: |-
-    {% if RAG %}To help you generate new {{targets}}, we have prepared the following information for you:
-    {{ RAG }}{% endif %}
+    {% if RAG %}
+    To assist you in generating new {{targets}}, we have provided the following information: {{RAG}}.
+    **Note:** The provided RAG is for reference only. 
+    You must carefully assess whether the RAG aligns with the {{targets}}. 
+    If it does not, it should not be used. Exercise caution and make your own judgment.
+    {% endif %}
     Also generate the relevant keys for the reasoning and the distilled knowledge that follows. For those keys, in particular for knowledge, explain in the context of the specific scenario to build up domain knowledge in the specific field rather than general knowledge.
 
 hypothesis2experiment:
@@ -35,8 +40,4 @@ hypothesis2experiment:
     {{ target_hypothesis }}
     The former hypothesis and the corresponding feedbacks are as follows:
     {{ hypothesis_and_feedback }}
-    The former proposed {{targets}} on similar hypothesis are as follows:
-    {{ target_list }}
-    To help you generate new {{targets}}, we have prepared the following information for you:
-    {{ RAG }}
     Please generate the new {{targets}} based on the information above.

rdagent/core/knowledge_base.py

@@ -14,12 +14,12 @@ class KnowledgeBase:
         if self.path is not None and self.path.exists():
             with self.path.open("rb") as f:
                 self.__dict__.update(
-                    pickle.load(f),
+                    pickle.load(f).__dict__,
                 )  # TODO: because we need to align with init function, we need a less hacky way to do this
 
     def dump(self) -> None:
         if self.path is not None:
             self.path.parent.mkdir(parents=True, exist_ok=True)
-            pickle.dump(self.__dict__, self.path.open("wb"))
+            pickle.dump(self, self.path.open("wb"))
         else:
             logger.warning("KnowledgeBase path is not set, dump failed.")

rdagent/scenarios/kaggle/experiment/playground-series-s3e11_template/fea_share_preprocess.py

@@ -23,6 +23,7 @@ def preprocess_script():
 
     # train
     train = pd.read_csv("/kaggle/input/train.csv")
+    train = train.drop(["id"], axis=1)
     train["store_sqft"] = train["store_sqft"].astype("category")
     train["salad"] = (train["salad_bar"] + train["prepared_food"]) / 2
     train["log_cost"] = np.log1p(train["cost"])

rdagent/scenarios/kaggle/experiment/scenario.py

@@ -20,6 +20,17 @@ from rdagent.scenarios.kaggle.knowledge_management.vector_base import (
 
 prompt_dict = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
 
+KG_ACTION_FEATURE_PROCESSING = "Feature processing"
+KG_ACTION_FEATURE_ENGINEERING = "Feature engineering"
+KG_ACTION_MODEL_FEATURE_SELECTION = "Model feature selection"
+KG_ACTION_MODEL_TUNING = "Model tuning"
+KG_ACTION_LIST = [
+    KG_ACTION_FEATURE_PROCESSING,
+    KG_ACTION_FEATURE_ENGINEERING,
+    KG_ACTION_MODEL_FEATURE_SELECTION,
+    KG_ACTION_MODEL_TUNING,
+]
+
 
 class KGScenario(Scenario):
     def __init__(self, competition: str) -> None:
@@ -54,6 +65,13 @@ class KGScenario(Scenario):
         self._simulator = self.simulator
         self._background = self.background
 
+        self.action_counts = dict.fromkeys(KG_ACTION_LIST, 0)
+        self.reward_estimates = {action: 0.0 for action in KG_ACTION_LIST}
+        self.reward_estimates["Model feature selection"] = 0.2
+        self.reward_estimates["Model tuning"] = 1.0
+        self.confidence_parameter = 1.0
+        self.initial_performance = 0.0
+
     def _analysis_competition_description(self):
         sys_prompt = (
             Environment(undefined=StrictUndefined)

rdagent/scenarios/kaggle/knowledge_management/prompts.yaml

@@ -6,7 +6,7 @@ extract_kaggle_knowledge_prompts:
     
     Please provide the analysis in the following JSON format:
     {
-      "content": "all provided content",
+      "content": "Put the provided content here",
       "title": "extracted title, if available",
       "competition_name": "extracted competition name",
       "task_category": "extracted task type, e.g., Classification, Regression",
@@ -80,4 +80,14 @@ extract_knowledge_graph_from_document:
     If you find no valuable insights in the document, please return an empty dict.
   
   user: |-
-    Document content: {{ document_content }}
+    Document content: {{ document_content }}
+
+refine_with_LLM:
+  system: |-
+    You are an experienced data science expert and an assistant, helping the user evaluate and improve content.
+  
+  user: |-
+    Here is the target: {{ target }}. 
+    Please evaluate whether the following RAG query result aligns with the target. 
+    If it does not, simply respond with "There are no relevant RAG results to support."
+    RAG query result: {{ text }}.

rdagent/scenarios/kaggle/knowledge_management/vector_base.py

@@ -4,8 +4,10 @@ from typing import List, Union
 
 import pandas as pd
 from _pytest.cacheprovider import json
+from jinja2 import Environment, StrictUndefined
 
 from rdagent.components.knowledge_management.vector_base import Document, PDVectorBase
+from rdagent.core.prompts import Prompts
 from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.kaggle.knowledge_management.extract_knowledge import (
@@ -225,12 +227,14 @@ class KaggleExperienceBase(PDVectorBase):
             document.create_embedding()
             self.add(document)
 
-    def search_experience(self, query: str, topk_k: int = 5, similarity_threshold: float = 0.1):
+    def search_experience(self, target: str, query: str, topk_k: int = 5, similarity_threshold: float = 0.1):
         """
-        Search for Kaggle experience posts related to the query
+        Search for Kaggle experience posts related to the query, initially filtered by the target.
 
         Parameters:
         ----------
+        target: str
+            The target context to refine the search query.
         query: str
             The search query to find relevant experience posts.
         topk_k: int, optional
@@ -243,15 +247,49 @@ class KaggleExperienceBase(PDVectorBase):
         List[KGKnowledgeMetaData], List[float]:
             A list of the most relevant documents and their similarities.
         """
-        search_results, similarities = super().search(query, topk_k=topk_k, similarity_threshold=similarity_threshold)
 
+        # Modify the query to include the target
+        modified_query = f"The target is {target}. And I need you to query {query} based on the {target}."
+
+        # First, search based on the modified query
+        search_results, similarities = super().search(
+            modified_query, topk_k=topk_k, similarity_threshold=similarity_threshold
+        )
+
+        # If the results do not match the target well, refine the search using LLM or further adjustment
         kaggle_docs = []
         for result in search_results:
             kg_doc = KGKnowledgeDocument().from_dict(result.__dict__)
+
+            gpt_feedback = self.refine_with_LLM(target, kg_doc)
+            if gpt_feedback:
+                kg_doc.content = gpt_feedback
+
             kaggle_docs.append(kg_doc)
 
         return kaggle_docs, similarities
 
+    def refine_with_LLM(self, target: str, text: str) -> str:
+        prompt_dict = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
+
+        sys_prompt = (
+            Environment(undefined=StrictUndefined).from_string(prompt_dict["refine_with_LLM"]["system"]).render()
+        )
+
+        user_prompt = (
+            Environment(undefined=StrictUndefined)
+            .from_string(prompt_dict["refine_with_LLM"]["user"])
+            .render(target=target, text=text)
+        )
+
+        response = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=sys_prompt,
+            json_mode=False,
+        )
+
+        return response
+
     def save(self, vector_df_path: Union[str, Path]):
         """
         Save the vector DataFrame to a file

rdagent/scenarios/kaggle/prompts.yaml

@@ -47,79 +47,80 @@ hypothesis_output_format: |-
     "concise_knowledge": "One line summary. Transferable knowledge based on theoretical principles. Use conditional grammar. eg. "If...., ..; When..., .; and etc" Make sure that you state things clearly without ambiguity. Eg. avoid saying "previous hypothesis", because one wouldn't know what that is."
   }
 
-hypothesis_specification: |-
-  There are different types of hypothesis that correspond to different types of actions. The specifications are quite important here:
-  1) feature_engineering:
-  description: We engineer the features for the sake of best model performance on the basis of engineering the most influential features.
-  type_of_feature_and_data_characteristics:
-    - Clearly define the feature type being introduced.
-    - Highlight the specific data patterns or characteristics the feature captures.
-    - Keep it focused—omit unnecessary details.
-  start_with_simple_features:
-    - Begin with straightforward and impactful features.
-    - Briefly explain why these features are expected to work.
-    - Avoid combining complex features at the outset.
-  increase_complexity_gradually:
-    - Add more complex features only after gathering experimental results.
-    - Discuss potential advantages and the trade-offs involved.
-    - Combine features only after simpler ones are tested and validated.
-  new_directions_and_optimizations:
-    - Justify any new direction based on data analysis or domain knowledge.
-    - Focus on one new direction at a time for clarity.
-    - If a hypothesis shows optimization potential (even without surpassing previous best results), explain why and proceed.
-  feature_library_and_task_management:
-    - Include features that improve performance in the feature library.
-    - Each generation should focus on 1-3 feature tasks, balancing simplicity with complexity.
+hypothesis_specification:
+  Feature engineering: |-
+    Action: Feature engineering
+    description: We engineer the features for the sake of best model performance on the basis of engineering the most influential features.
+    type_of_feature_and_data_characteristics:
+      - Clearly define the feature type being introduced.
+      - Highlight the specific data patterns or characteristics the feature captures.
+      - Keep it focused—omit unnecessary details.
+    start_with_simple_features:
+      - Begin with straightforward and impactful features.
+      - Briefly explain why these features are expected to work.
+      - Avoid combining complex features at the outset.
+    increase_complexity_gradually:
+      - Add more complex features only after gathering experimental results.
+      - Discuss potential advantages and the trade-offs involved.
+      - Combine features only after simpler ones are tested and validated.
+    new_directions_and_optimizations:
+      - Justify any new direction based on data analysis or domain knowledge.
+      - Focus on one new direction at a time for clarity.
+      - If a hypothesis shows optimization potential (even without surpassing previous best results), explain why and proceed.
+    feature_library_and_task_management:
+      - Include features that improve performance in the feature library.
+      - Each generation should focus on 1-3 feature tasks, balancing simplicity with complexity.
 
-  2) feature processing:
-  define_the_processing_method:
-    - Clearly state the type of feature processing.
-    - Explain how this processing captures data patterns or improves feature usefulness.
-    - Avoid redundant details.
-  begin_with_simple_processing:
-    - Start with simple, effective processing methods.
-    - Concisely explain why these methods should improve model performance.
-    - Introduce complex processing only after gathering experimental results.
-  introduce_complexity_gradually:
-    - Add more sophisticated processing methods step-by-step, after validation.
-    - Discuss the advantages, challenges, and trade-offs of advanced processing.
-    - Validate simpler methods before combining them with complex ones.
+  Feature processing: |-
+    Action: Feature processing
+    Define_the_processing_method:
+      - Clearly state the type of feature processing.
+      - Explain how this processing captures data patterns or improves feature usefulness.
+      - Avoid redundant details.
+    Begin_with_simple_processing:
+      - Start with simple, effective processing methods.
+      - Concisely explain why these methods should improve model performance.
+      - Introduce complex processing only after gathering experimental results.
+    Introduce_complexity_gradually:
+      - Add more sophisticated processing methods step-by-step, after validation.
+      - Discuss the advantages, challenges, and trade-offs of advanced processing.
+      - Validate simpler methods before combining them with complex ones.
 
-  3) model feature selection:
-  selection_based_on_model_type:
-    - Specify which features are being selected and explain why, considering the model type (e.g., NN, Random Forest, LightGBM, XGBoost).
-    - Ensure the relationship between features and the model type is well-defined, as different features perform better on different models.
-  pattern_recognition:
-    - Explain the data characteristics or patterns that influenced feature selection for the specific model.
-    - Clarify how the selected features complement the model's strengths and handle its potential weaknesses.
+  Model feature selection: |-
+    Selection_based_on_model_type:
+      - Specify which features are being selected and explain why, considering the model type (e.g., NN, Random Forest, LightGBM, XGBoost).
+      - Ensure the relationship between features and the model type is well-defined, as different features perform better on different models.
+    Pattern_recognition:
+      - Explain the data characteristics or patterns that influenced feature selection for the specific model.
+      - Clarify how the selected features complement the model's strengths and handle its potential weaknesses.
 
-  4) model_design_and_tuning:
-  Explain the hypothesis clearly with valuable information. What kind of model are you building/tuning? What do you think is true? How you are revising and why? What are some innvations? 
-  focus_on_architecture_or_hyper_parameter_tuning_or_both:
-    - Focus on designing new model architectures one at a time OR hyper-parameter tuning OR both.
-    - Each hypothesis should introduce a novel architecture or a significant modification to an existing one, while leveraging previous experiences and the hypothesis history.
-    - Optimize one model at a time, iterating until its potential is fully explored. Switch to a new model only when you believe the current model’s potential has been exhausted.
-  specific_to_model_type:
-    - Note that any types of tuning or model design must be specific to the model types available in our workspace.
-    - Clearly define the model type (e.g., Neural Network Models (eg, MLP, CNN, RNN, LSTM, GRU etc.), XGBoost, RandomForest, LightGBM) and the architecture/tuning being introduced.
-    - Ensure the architecture or tuning aligns with the data characteristics and the strengths or limitations of the specific model.
-  rationale_behind_architecture_and_tuning:
-    - Explain the innovation or reasoning behind the architectural design or tuning approach.
-    - Justify how the new structure or parameter change captures data patterns more effectively, improves learning efficiency, or enhances predictive power.
-  start_simple_innovate_gradually:
-    - Start with innovative yet simple changes to ensure each iteration is well-tested and the results are well-understood.
-    - Gradually introduce more complex architectural changes or hyper-parameter adjustments based on gathered results and insights.
-  introduce_one_innovation_at_a_time:
-    - Focus on testing one key innovation at a time to isolate its impact on performance.
-    - Avoid combining multiple innovations in a single iteration to maintain clarity in performance results.
-  balance_innovation_with_performance:
-    - Strive for a balance between creative design and practical, effective performance.
-    - If a design or tuning shows strong performance, document it in a "library" for future iterations.
-  iterative_testing_and_refinement:
-    - After each test, evaluate and refine the model architecture or tuning based on observed performance and data patterns.
-    - If a hypothesis shows potential but doesn't surpass previous results, continue optimizing in that direction.
-  hypothesis_statement:
-    - For each hypothesis, specify the exact innovation or tuning approach and explain why it's expected to enhance performance for the chosen model type.
+  Model tuning: |-
+    Explain the hypothesis clearly with valuable information. What kind of model are you building/tuning? What do you think is true? How you are revising and why? What are some innvations? 
+    Focus_on_architecture_or_hyper_parameter_tuning_or_both:
+      - Focus on designing new model architectures one at a time OR hyper-parameter tuning OR both.
+      - Each hypothesis should introduce a novel architecture or a significant modification to an existing one, while leveraging previous experiences and the hypothesis history.
+      - Optimize one model at a time, iterating until its potential is fully explored. Switch to a new model only when you believe the current model’s potential has been exhausted.
+    Specific_to_model_type:
+      - Note that any types of tuning or model design must be specific to the model types available in our workspace.
+      - Clearly define the model type (e.g., Neural Network Models (eg, MLP, CNN, RNN, LSTM, GRU etc.), XGBoost, RandomForest, LightGBM) and the architecture/tuning being introduced.
+      - Ensure the architecture or tuning aligns with the data characteristics and the strengths or limitations of the specific model.
+    Rationale_behind_architecture_and_tuning:
+      - Explain the innovation or reasoning behind the architectural design or tuning approach.
+      - Justify how the new structure or parameter change captures data patterns more effectively, improves learning efficiency, or enhances predictive power.
+    Start_simple_innovate_gradually:
+      - Start with innovative yet simple changes to ensure each iteration is well-tested and the results are well-understood.
+      - Gradually introduce more complex architectural changes or hyper-parameter adjustments based on gathered results and insights.
+    Introduce_one_innovation_at_a_time:
+      - Focus on testing one key innovation at a time to isolate its impact on performance.
+      - Avoid combining multiple innovations in a single iteration to maintain clarity in performance results.
+    Balance_innovation_with_performance:
+      - Strive for a balance between creative design and practical, effective performance.
+      - If a design or tuning shows strong performance, document it in a "library" for future iterations.
+    Iterative_testing_and_refinement:
+      - After each test, evaluate and refine the model architecture or tuning based on observed performance and data patterns.
+      - If a hypothesis shows potential but doesn't surpass previous results, continue optimizing in that direction.
+    Hypothesis_statement:
+      - For each hypothesis, specify the exact innovation or tuning approach and explain why it's expected to enhance performance for the chosen model type.
 
 
 feature_experiment_output_format: |-
@@ -158,7 +159,8 @@ model_experiment_output_format: |-
           "hyperparameter_name_2": "value of hyperparameter 2",
           "hyperparameter_name_3": "value of hyperparameter 3"
       },
-      "model_type": "Select only one model type: XGBoost, RandomForest, LightGBM, or NN. The primary model must be unique, but you may use auxiliary models for support if you think it can have a good result like choosing A model as the main model, with B Model used for auxiliary support or optimization on specific details."
+      "model_type": "Please select only **one** model type from the following four options: XGBoost, RandomForest, LightGBM, or NN. The selected model must be unique and used as the **primary model**. You may choose an auxiliary model for support or optimization on specific tasks if necessary, but the primary model must come from the provided options."
+
   }
   Usually, a larger model works better than a smaller one. Hence, the parameters should be larger.
 

rdagent/scenarios/kaggle/proposal/proposal.py

@@ -14,9 +14,11 @@ from rdagent.components.proposal.model_proposal import (
     ModelHypothesis2Experiment,
     ModelHypothesisGen,
 )
+from rdagent.core.exception import ModelEmptyError
 from rdagent.core.prompts import Prompts
 from rdagent.core.proposal import Hypothesis, Scenario, Trace
 from rdagent.scenarios.kaggle.experiment.kaggle_experiment import (
+    KG_SELECT_MAPPING,
     KGFactorExperiment,
     KGModelExperiment,
 )
@@ -29,16 +31,13 @@ from rdagent.scenarios.kaggle.knowledge_management.vector_base import (
 prompt_dict = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
 
 
-KG_ACTION_FEATURE_PROCESSING = "Feature processing"
-KG_ACTION_FEATURE_ENGINEERING = "Feature engineering"
-KG_ACTION_MODEL_FEATURE_SELECTION = "Model feature selection"
-KG_ACTION_MODEL_TUNING = "Model tuning"
-KG_ACTION_LIST = [
-    KG_ACTION_FEATURE_PROCESSING,
+from rdagent.scenarios.kaggle.experiment.scenario import (
     KG_ACTION_FEATURE_ENGINEERING,
+    KG_ACTION_FEATURE_PROCESSING,
+    KG_ACTION_LIST,
     KG_ACTION_MODEL_FEATURE_SELECTION,
     KG_ACTION_MODEL_TUNING,
-]
+)
 
 
 class KGHypothesis(Hypothesis):
@@ -82,19 +81,13 @@ class KGHypothesisGen(ModelHypothesisGen):
 
     def __init__(self, scen: Scenario) -> Tuple[dict, bool]:
         super().__init__(scen)
-        self.action_counts = dict.fromkeys(KG_ACTION_LIST, 0)
-        self.reward_estimates = {action: 0.0 for action in KG_ACTION_LIST}
-        self.reward_estimates["Model feature selection"] = 0.2
-        self.reward_estimates["Model tuning"] = 1.0
-        self.confidence_parameter = 1.0
-        self.initial_performance = 0.0
 
-    def generate_RAG_content(self, trace: Trace, hypothesis_and_feedback: str) -> str:
+    def generate_RAG_content(self, trace: Trace, hypothesis_and_feedback: str, target: str = None) -> str:
         if self.scen.if_using_vector_rag:
             if self.scen.mini_case:
-                rag_results, _ = self.scen.vector_base.search_experience(hypothesis_and_feedback, topk_k=1)
+                rag_results, _ = self.scen.vector_base.search_experience(target, hypothesis_and_feedback, topk_k=1)
             else:
-                rag_results, _ = self.scen.vector_base.search_experience(hypothesis_and_feedback, topk_k=5)
+                rag_results, _ = self.scen.vector_base.search_experience(target, hypothesis_and_feedback, topk_k=5)
             return "\n".join([doc.content for doc in rag_results])
         if self.scen.if_using_graph_rag is False or trace.knowledge_base is None:
             return None
@@ -188,37 +181,42 @@ class KGHypothesisGen(ModelHypothesisGen):
                 prev_result = prev_entry[1].result
                 performance_t_minus_1 = prev_result.get("performance", 0.0)
             else:
-                performance_t_minus_1 = self.initial_performance
+                performance_t_minus_1 = self.scen.initial_performance
+
+            if self.scen.evaluation_metric_direction:
+                reward = (performance_t - performance_t_minus_1) / max(performance_t_minus_1, 1e-8)
+            else:
+                reward = (performance_t_minus_1 - performance_t) / max(performance_t_minus_1, 1e-8)
 
             reward = (performance_t - performance_t_minus_1) / performance_t_minus_1
-            n_o = self.action_counts[last_action]
-            mu_o = self.reward_estimates[last_action]
-            self.reward_estimates[last_action] += (reward - mu_o) / n_o
+            n_o = self.scen.action_counts[last_action]
+            mu_o = self.scen.reward_estimates[last_action]
+            self.scen.scen.reward_estimates[last_action] += (reward - mu_o) / n_o
         else:
             # First iteration, nothing to update
             pass
 
     def execute_next_action(self, trace: Trace) -> str:
-        actions = list(self.action_counts.keys())
-        t = sum(self.action_counts.values()) + 1
+        actions = list(self.scen.action_counts.keys())
+        t = sum(self.scen.action_counts.values()) + 1
 
         # If any action has not been tried yet, select it
         for action in actions:
-            if self.action_counts[action] == 0:
+            if self.scen.action_counts[action] == 0:
                 selected_action = action
-                self.action_counts[selected_action] += 1
+                self.scen.action_counts[selected_action] += 1
                 return selected_action
 
-        c = self.confidence_parameter
+        c = self.scen.confidence_parameter
         ucb_values = {}
         for action in actions:
-            mu_o = self.reward_estimates[action]
-            n_o = self.action_counts[action]
+            mu_o = self.scen.reward_estimates[action]
+            n_o = self.scen.action_counts[action]
             ucb = mu_o + c * math.sqrt(math.log(t) / n_o)
             ucb_values[action] = ucb
         # Select action with highest UCB
         selected_action = max(ucb_values, key=ucb_values.get)
-        self.action_counts[selected_action] += 1
+        self.scen.action_counts[selected_action] += 1
         return selected_action
 
     def prepare_context(self, trace: Trace) -> Tuple[dict, bool]:
@@ -237,11 +235,15 @@ class KGHypothesisGen(ModelHypothesisGen):
 
         context_dict = {
             "hypothesis_and_feedback": hypothesis_and_feedback,
-            "RAG": self.generate_RAG_content(trace, hypothesis_and_feedback),
+            "RAG": self.generate_RAG_content(
+                trace=trace,
+                hypothesis_and_feedback=hypothesis_and_feedback,
+                target=action if self.scen.if_action_choosing_based_on_UCB else None,
+            ),
             "hypothesis_output_format": prompt_dict["hypothesis_output_format"],
             "hypothesis_specification": (
                 f"next experiment action is {action}" if self.scen.if_action_choosing_based_on_UCB else None,
-                prompt_dict["hypothesis_specification"],
+                prompt_dict["hypothesis_specification"][action],
             ),
         }
         return context_dict, True
@@ -287,7 +289,8 @@ class KGHypothesis2Experiment(ModelHypothesis2Experiment):
 
         model_list = []
         for experiment in experiment_list:
-            model_list.extend(experiment.sub_tasks)
+            for sub_task in experiment.sub_tasks:
+                model_list.extend(sub_task.get_task_information())
 
         return {
             "target_hypothesis": str(hypothesis),
@@ -328,6 +331,11 @@ class KGHypothesis2Experiment(ModelHypothesis2Experiment):
     def convert_model_experiment(self, response: str, trace: Trace) -> KGModelExperiment:
         response_dict = json.loads(response)
         tasks = []
+        model_type = response_dict.get("model_type", "Model type not provided")
+        if model_type not in KG_SELECT_MAPPING:
+            raise ModelEmptyError(
+                f"Invalid model type '{model_type}'. Allowed model types are: {', '.join(KG_SELECT_MAPPING)}."
+            )
         tasks.append(
             ModelTask(
                 name=response_dict.get("model_name", "Model name not provided"),