feat: enhance timeout management and knowledge base handling in CoSTEER components (#1130)

* feat: enhance timeout management and knowledge base handling in CoSTEER components

* fix a little bug

* fix small bug

* fix a small bug

* Update rdagent/scenarios/data_science/loop.py

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>

* add scale check

* fix a small bug

* fix CI

* use dynamic chat_token_limit & remove repeated lines

* fix CI

* remove useless comment

* fix small bug

* update draft appendix

* fix prompt

* add code correctness as top priority

---------

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
This commit is contained in:
Xu Yang
2025-07-31 13:06:28 +08:00
committed by GitHub
parent ce25bafbf7
commit 1aa5cc2535
26 changed files with 337 additions and 131 deletions
+4 -2
View File
@@ -119,12 +119,14 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
model_architecture_suggestion_time_percent: float = 0.75
allow_longer_timeout: bool = False
coder_longer_timeout_multiplier: int = 3
runner_longer_timeout_multiplier: int = 2
coder_longer_timeout_multiplier_upper: int = 3
runner_longer_timeout_multiplier_upper: int = 2
timeout_increase_stage: float = 0.3
#### hypothesis critique and rewrite
enable_hypo_critique_rewrite: bool = True
"""Enable hypothesis critique and rewrite stages for improving hypothesis quality"""
enable_scale_check: bool = False
DS_RD_SETTING = DataScienceBasePropSetting()
+19 -37
View File
@@ -6,13 +6,11 @@ from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
from rdagent.components.coder.CoSTEER.evaluators import CoSTEERMultiFeedback
from rdagent.components.coder.CoSTEER.evolvable_subjects import EvolvingItem
from rdagent.components.coder.CoSTEER.knowledge_management import (
CoSTEERKnowledgeBaseV1,
CoSTEERKnowledgeBaseV2,
CoSTEERRAGStrategyV1,
CoSTEERRAGStrategyV2,
)
from rdagent.core.developer import Developer
from rdagent.core.evaluation import Evaluator, Feedback
from rdagent.core.evaluation import Evaluator
from rdagent.core.evolving_agent import EvolvingStrategy, RAGEvoAgent
from rdagent.core.exception import CoderError
from rdagent.core.experiment import Experiment
@@ -26,8 +24,8 @@ class CoSTEER(Developer[Experiment]):
settings: CoSTEERSettings,
eva: Evaluator,
es: EvolvingStrategy,
evolving_version: int,
*args,
evolving_version: int = 2,
max_seconds: int | None = None,
with_knowledge: bool = True,
with_feedback: bool = True,
@@ -37,6 +35,8 @@ class CoSTEER(Developer[Experiment]):
**kwargs,
) -> None:
super().__init__(*args, **kwargs)
self.settings = settings
self.max_loop = settings.max_loop if max_loop is None else max_loop
self.max_seconds = max_seconds
self.knowledge_base_path = (
@@ -54,37 +54,22 @@ class CoSTEER(Developer[Experiment]):
self.evaluator = eva
self.evolving_version = evolving_version
# init knowledge base
self.knowledge_base = self.load_or_init_knowledge_base(
former_knowledge_base_path=self.knowledge_base_path,
component_init_list=[],
)
# init rag method
self.rag = (
CoSTEERRAGStrategyV2(self.knowledge_base, settings=settings)
if self.evolving_version == 2
else CoSTEERRAGStrategyV1(self.knowledge_base, settings=settings)
)
def load_or_init_knowledge_base(self, former_knowledge_base_path: Path = None, component_init_list: list = []):
if former_knowledge_base_path is not None and former_knowledge_base_path.exists():
knowledge_base = pickle.load(open(former_knowledge_base_path, "rb"))
if self.evolving_version == 1 and not isinstance(knowledge_base, CoSTEERKnowledgeBaseV1):
raise ValueError("The former knowledge base is not compatible with the current version")
elif self.evolving_version == 2 and not isinstance(
knowledge_base,
CoSTEERKnowledgeBaseV2,
):
raise ValueError("The former knowledge base is not compatible with the current version")
else:
knowledge_base = (
CoSTEERKnowledgeBaseV2(
init_component_list=component_init_list,
)
if self.evolving_version == 2
else CoSTEERKnowledgeBaseV1()
CoSTEERRAGStrategyV2(
settings=settings,
former_knowledge_base_path=self.knowledge_base_path,
dump_knowledge_base_path=self.new_knowledge_base_path,
evolving_version=self.evolving_version,
)
return knowledge_base
if self.evolving_version == 2
else CoSTEERRAGStrategyV1(
settings=settings,
former_knowledge_base_path=self.knowledge_base_path,
dump_knowledge_base_path=self.new_knowledge_base_path,
evolving_version=self.evolving_version,
)
)
def develop(self, exp: Experiment) -> Experiment:
@@ -98,6 +83,8 @@ class CoSTEER(Developer[Experiment]):
with_knowledge=self.with_knowledge,
with_feedback=self.with_feedback,
knowledge_self_gen=self.knowledge_self_gen,
enable_filelock=self.settings.enable_filelock,
filelock_path=self.settings.filelock_path,
)
start_datetime = datetime.now()
@@ -116,11 +103,6 @@ class CoSTEER(Developer[Experiment]):
if self.with_feedback and self.filter_final_evo:
evo_exp = self._exp_postprocess_by_feedback(evo_exp, self.evolve_agent.evolving_trace[-1].feedback)
# save new knowledge base
if self.new_knowledge_base_path is not None:
with self.new_knowledge_base_path.open("wb") as f:
pickle.dump(self.knowledge_base, f)
logger.info(f"New knowledge base saved to {self.new_knowledge_base_path}")
exp.sub_workspace_list = evo_exp.sub_workspace_list
exp.experiment_workspace = evo_exp.experiment_workspace
return exp
@@ -33,6 +33,9 @@ class CoSTEERSettings(ExtendedBaseSettings):
new_knowledge_base_path: Union[str, None] = None
"""Path to the new knowledge base"""
enable_filelock: bool = False
filelock_path: Union[str, None] = None
max_seconds_multiplier: int = 10**6
@@ -75,6 +75,8 @@ class MultiProcessEvolvingStrategy(EvolvingStrategy):
evolving_trace: list[EvoStep] = [],
**kwargs,
) -> EvolvingItem:
code_list = [None for _ in range(len(evo.sub_tasks))]
# 1.找出需要evolve的task
to_be_finished_task_index: list[int] = []
for index, target_task in enumerate(evo.sub_tasks):
@@ -82,9 +84,9 @@ class MultiProcessEvolvingStrategy(EvolvingStrategy):
if target_task_desc in queried_knowledge.success_task_to_knowledge_dict:
# NOTE: very weird logic:
# it depends on the knowledge to set the already finished task
evo.sub_workspace_list[index] = queried_knowledge.success_task_to_knowledge_dict[
code_list[index] = queried_knowledge.success_task_to_knowledge_dict[
target_task_desc
].implementation
].implementation.file_dict
elif (
target_task_desc not in queried_knowledge.success_task_to_knowledge_dict
and target_task_desc not in queried_knowledge.failed_task_info_set
@@ -111,7 +113,6 @@ class MultiProcessEvolvingStrategy(EvolvingStrategy):
],
n=RD_AGENT_SETTINGS.multi_proc_n,
)
code_list = [None for _ in range(len(evo.sub_tasks))]
for index, target_index in enumerate(to_be_finished_task_index):
code_list[target_index] = result[index]
@@ -2,6 +2,7 @@ from __future__ import annotations
import copy
import json
import pickle
import random
import re
from itertools import combinations
@@ -51,6 +52,53 @@ class CoSTEERKnowledge(Knowledge):
"""
class CoSTEERRAGStrategy(RAGStrategy):
def __init__(self, *args, dump_knowledge_base_path: Path = None, **kwargs):
super().__init__(*args, **kwargs)
self.dump_knowledge_base_path = dump_knowledge_base_path
def load_or_init_knowledge_base(
self, former_knowledge_base_path: Path = None, component_init_list: list = [], evolving_version: int = 2
) -> EvolvingKnowledgeBase:
if former_knowledge_base_path is not None and former_knowledge_base_path.exists():
knowledge_base = pickle.load(open(former_knowledge_base_path, "rb"))
if evolving_version == 1 and not isinstance(knowledge_base, CoSTEERKnowledgeBaseV1):
raise ValueError("The former knowledge base is not compatible with the current version")
elif evolving_version == 2 and not isinstance(
knowledge_base,
CoSTEERKnowledgeBaseV2,
):
raise ValueError("The former knowledge base is not compatible with the current version")
else:
knowledge_base = (
CoSTEERKnowledgeBaseV2(
init_component_list=component_init_list,
)
if evolving_version == 2
else CoSTEERKnowledgeBaseV1()
)
return knowledge_base
def dump_knowledge_base(self):
if self.dump_knowledge_base_path is None:
logger.warning("Dump knowledge base path is not set, skip dumping.")
else:
if not self.dump_knowledge_base_path.parent.exists():
self.dump_knowledge_base_path.parent.mkdir(parents=True, exist_ok=True)
with open(self.dump_knowledge_base_path, "wb") as f:
pickle.dump(self.knowledgebase, f)
def load_dumped_knowledge_base(self, *args, **kwargs):
if self.dump_knowledge_base_path is None:
logger.warning("Dump knowledge base path is not set, skip dumping.")
elif not Path(self.dump_knowledge_base_path).exists():
logger.info(f"Dumped knowledge base {self.dump_knowledge_base_path} does not exist, skip loading.")
else:
with open(self.dump_knowledge_base_path, "rb") as f:
self.knowledgebase = pickle.load(f)
logger.info(f"Loaded dumped knowledge base from {self.dump_knowledge_base_path}")
class CoSTEERQueriedKnowledge(QueriedKnowledge):
def __init__(self, success_task_to_knowledge_dict: dict = {}, failed_task_info_set: set = set()) -> None:
self.success_task_to_knowledge_dict = success_task_to_knowledge_dict
@@ -85,9 +133,9 @@ class CoSTEERQueriedKnowledgeV1(CoSTEERQueriedKnowledge):
super().__init__(*args, **kwargs)
class CoSTEERRAGStrategyV1(RAGStrategy):
def __init__(self, knowledgebase: CoSTEERKnowledgeBaseV1, settings: CoSTEERSettings) -> None:
super().__init__(knowledgebase)
class CoSTEERRAGStrategyV1(CoSTEERRAGStrategy):
def __init__(self, settings: CoSTEERSettings, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.current_generated_trace_count = 0
self.settings = settings
@@ -213,9 +261,9 @@ class CoSTEERQueriedKnowledgeV2(CoSTEERQueriedKnowledgeV1):
)
class CoSTEERRAGStrategyV2(RAGStrategy):
def __init__(self, knowledgebase: CoSTEERKnowledgeBaseV2, settings: CoSTEERSettings) -> None:
super().__init__(knowledgebase)
class CoSTEERRAGStrategyV2(CoSTEERRAGStrategy):
def __init__(self, settings: CoSTEERSettings, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.current_generated_trace_count = 0
self.settings = settings
@@ -249,6 +297,12 @@ class CoSTEERRAGStrategyV2(RAGStrategy):
target_task_information not in self.knowledgebase.success_task_to_knowledge_dict
and implementation is not None
):
if target_task_information not in self.knowledgebase.task_to_component_nodes:
self.knowledgebase.task_to_component_nodes[target_task_information] = (
self.analyze_component(
target_task_information,
)
)
self.knowledgebase.working_trace_knowledge.setdefault(target_task_information, []).append(
single_knowledge,
) # save to working trace
@@ -465,7 +519,6 @@ class CoSTEERRAGStrategyV2(RAGStrategy):
self.knowledgebase.task_to_component_nodes[target_task_information] = self.analyze_component(
target_task_information,
)
component_analysis_result = self.knowledgebase.task_to_component_nodes[target_task_information]
if len(component_analysis_result) > 1:
@@ -557,12 +610,14 @@ class CoSTEERRAGStrategyV2(RAGStrategy):
queried_from_gt_knowledge_list = [
knowledge
for knowledge in queried_knowledge_list
if knowledge.feedback is not None and knowledge.feedback.final_decision_based_on_gt == True
if knowledge.feedback is not None
and (
hasattr(knowledge.feedback, "final_decision_based_on_gt")
and knowledge.feedback.final_decision_based_on_gt == True
)
]
queried_without_gt_knowledge_list = [
knowledge
for knowledge in queried_knowledge_list
if knowledge.feedback is not None and knowledge.feedback.final_decision_based_on_gt == False
knowledge for knowledge in queried_knowledge_list if knowledge not in queried_from_gt_knowledge_list
]
queried_from_gt_knowledge_count = max(
min((v2_query_component_limit // 2 + 1), len(queried_from_gt_knowledge_list)),
@@ -63,11 +63,6 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
data_folder_info = self.scen.processed_data_folder_description
pipeline_task_info = target_task.get_task_information()
queried_similar_successful_knowledge = (
queried_knowledge.task_to_similar_task_successful_knowledge[pipeline_task_info]
if queried_knowledge is not None
else []
)
queried_former_failed_knowledge = (
queried_knowledge.task_to_former_failed_traces[pipeline_task_info] if queried_knowledge is not None else []
)
@@ -82,7 +77,6 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
system_prompt = T(".prompts:pipeline_coder.system").r(
task_desc=pipeline_task_info,
queried_similar_successful_knowledge=queried_similar_successful_knowledge,
queried_former_failed_knowledge=queried_former_failed_knowledge[0],
out_spec=PythonAgentOut.get_spec(),
runtime_environment=self.scen.get_runtime_environment(),
@@ -158,10 +158,17 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
else:
eda_output = implementation.file_dict.get("EDA.md", None)
queried_similar_successful_knowledge = (
queried_knowledge.task_to_similar_task_successful_knowledge[target_task.get_task_information()]
if queried_knowledge is not None
else []
)
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,
mle_check=(DS_RD_SETTING.sample_data_by_LLM and test_eval.is_sub_enabled(self.scen.competition)),
queried_similar_successful_knowledge=queried_similar_successful_knowledge,
)
user_prompt = T(".prompts:pipeline_eval.user").r(
scenario=self.scen.get_scenario_all_desc(eda_output=eda_output),
@@ -28,19 +28,6 @@ pipeline_coder:
# Specification your code should follow
{% include "scenarios.data_science.share:component_spec.Pipeline" %}
{% if queried_similar_successful_knowledge|length != 0 or queried_former_failed_knowledge|length != 0 %}
# Relevant Information for This Task
{% endif %}
{% if queried_similar_successful_knowledge|length != 0 %}
## Successful Implementations for Similar Models
====={% for similar_successful_knowledge in queried_similar_successful_knowledge %} Model {{ loop.index }}:=====
{{ similar_successful_knowledge.target_task.get_task_information() }}
=====Code:=====
{{ similar_successful_knowledge.implementation.all_codes }}
{% endfor %}
{% endif %}
{% if queried_former_failed_knowledge|length != 0 %}
## Previous Failed Attempts
{% for former_failed_knowledge in queried_former_failed_knowledge %} Attempt {{ loop.index }}:
@@ -102,6 +89,16 @@ pipeline_coder:
```bash
python main.py --debug
```
Please simulate the following code to check whether the code is running in debug mode:
```python
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--debug', action='store_true', help='Run in debug mode')
args = parser.parse_args()
DEBUG = False
if args.debug:
DEBUG = True
```
In debug mode, you should only sample ten percent of the training data and run the minimum epochs to quickly test the correctness of the code.
In debug mode, you should implement a timer to measure the time taken for your debug configuration and estimate the time required for the full run. Your timer should only measure the time taken for the training part, not the data loading or feature engineering part.
For example:
@@ -115,7 +112,8 @@ pipeline_coder:
```
In debug mode, your code should run faster, so the environment will set a shorter time limit than the standard time limit for your code.
For example, you can sample ten percent of the training data and run for one epoch, then the full run with ten epochs will take one hundred times the time taken for the debug run. The scale is calculated by yourself depending on the data sampling and epoch number you choose. If your full run enables early stopping, the scale should be smaller considering the early stopping will stop the training earlier than the full epochs.
You should sample the data after train valid split. When you split the data after sampling, you might get a class with only one sample which might cause the split strategy to fail.
Be careful about the train-valid split strategy. StratifiedShuffleSplit is highly risk since the data has some categories with only one sample. If you use StratifiedShuffleSplit, you should consider using a try-except block to catch the error and use a different split strategy if the error occurs.
You should sample the data after train valid split. When you split the data after sampling, you might get a class with only one sample which might cause the split strategy to fail.
Your debug code should run exactly the same as the full run, except for the data sampling and epoch number, to ensure the correctness of the code.
You should print total time and estimated time in standard output using print function in the following schema:
=== Start of Debug Information ===
@@ -135,16 +133,17 @@ pipeline_coder:
{% endif %}
## General Guidelines
1. Use the print() function for all output; do not use the logging module.
2. **Avoid all hard-coded values (e.g., fixed dataset sizes)**. Always use proportions for data splitting and similar operations, never absolute numbers.
3. Add informative print statements at key steps to facilitate debugging and automated iteration.
4. For model training, use reasonable epoch numbers. ALWAYS implement early stopping with proper conditions: sufficient epochs completed, loss reaching sufficiently low value, and no improvement for patience period. Save best model checkpoints based on validation performance.
5. Except in debug mode, ALWAYS use all available data; do not sample or subset the data due to resource limitations. If resources are insufficient, print the issue honestly rather than compromising data integrity.
6. Do not use tqdm or similar progress bar tools.
7. **Try-except blocks are ONLY allowed when reading files. If no files are successfully read, it indicates incorrect file paths or reading methods, not a try-except issue. Try-except is PROHIBITED elsewhere in the code. Assert statements are PROHIBITED throughout the entire code.**
8. ATTENTION: ALWAYS use the best saved model (not necessarily final epoch) for predictions. **NEVER create dummy/placeholder submissions (e.g., all 1s, random values)**. If training fails, report failure honestly rather than generating fake submission files.
9. You should ALWAYS generate the complete code rather than partial code.
10. Strictly follow all specifications and general guidelines described above.
1. Code correctness is the top priority. Ensure your code is runnable and produces the expected output even if some task requirements are not fully met because the task itself might contain some errors like the wrong package name or wrong package function names.
2. Use the print() function for all output; do not use the logging module.
3. **Avoid all hard-coded values (e.g., fixed dataset sizes)**. Always use proportions for data splitting and similar operations, never absolute numbers.
4. Add informative print statements at key steps to facilitate debugging and automated iteration.
5. For model training, use reasonable epoch numbers. ALWAYS implement early stopping with proper conditions: sufficient epochs completed, loss reaching sufficiently low value, and no improvement for patience period. Save best model checkpoints based on validation performance.
6. Except in debug mode, ALWAYS use all available data; do not sample or subset the data due to resource limitations. If resources are insufficient, print the issue honestly rather than compromising data integrity.
7. Do not use tqdm or similar progress bar tools.
8. **Try-except blocks are ONLY allowed when reading files. If no files are successfully read, it indicates incorrect file paths or reading methods, not a try-except issue. Try-except is PROHIBITED elsewhere in the code. Assert statements are PROHIBITED throughout the entire code.**
9. ATTENTION: ALWAYS use the best saved model (not necessarily final epoch) for predictions. **NEVER create dummy/placeholder submissions (e.g., all 1s, random values)**. If training fails, report failure honestly rather than generating fake submission files.
10. You should ALWAYS generate the complete code rather than partial code.
11. Strictly follow all specifications and general guidelines described above.
### Output Format
{% if out_spec %}
@@ -269,6 +268,18 @@ pipeline_eval:
- If the submission check returns an error message, you should set the "final_decision" to false and clearly document the issues in the "return_checking" field.
{% endif %}
{% if queried_similar_successful_knowledge|length != 0 %}
### Step 6: Similar Successful Implementations to help Code Improvement
The user has done several similar tasks and get some successful implementations. These code might not be implemented to the same task, but they are similar to your task and they might work well on your dataset.
Please refer to these successful implementation and provide your suggestions in your response on how to correct your current code based on these successful implementations.
## Successful Implementations for Similar Tasks
====={% for similar_successful_knowledge in queried_similar_successful_knowledge %} Similar Task {{ loop.index }}:=====
{{ similar_successful_knowledge.target_task.get_task_information() }}
=====Code:=====
{{ similar_successful_knowledge.implementation.all_codes }}
{% endfor %}
{% endif %}
## Output Format
Please respond with your feedback in the following JSON format without anything else.
```json
@@ -103,7 +103,7 @@ class FactorCodeEvaluator(FactorEvaluator):
user_prompt=user_prompt,
system_prompt=system_prompt,
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
else:
@@ -509,7 +509,7 @@ class FactorFinalDecisionEvaluator(FactorEvaluator):
user_prompt=user_prompt,
system_prompt=system_prompt,
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
else:
@@ -45,7 +45,7 @@ class FactorMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
APIBackend().build_messages_and_calculate_token(
user_prompt=error_summary_user_prompt, system_prompt=error_summary_system_prompt
)
< LLM_SETTINGS.chat_token_limit
< APIBackend().chat_token_limit
):
break
elif len(queried_similar_error_knowledge_to_render) > 0:
@@ -124,7 +124,7 @@ class FactorMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
)
if (
APIBackend().build_messages_and_calculate_token(user_prompt=user_prompt, system_prompt=system_prompt)
< LLM_SETTINGS.chat_token_limit
< APIBackend().chat_token_limit
):
break
elif len(queried_former_failed_knowledge_to_render) > 1:
@@ -87,7 +87,7 @@ class ModelCodeEvaluator(CoSTEEREvaluator):
user_prompt=user_prompt,
system_prompt=system_prompt,
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
else:
@@ -142,7 +142,7 @@ class ModelFinalEvaluator(CoSTEEREvaluator):
user_prompt=user_prompt,
system_prompt=system_prompt,
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
else:
@@ -61,7 +61,7 @@ class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
user_prompt=user_prompt,
system_prompt=system_prompt,
)
< LLM_SETTINGS.chat_token_limit
< APIBackend().chat_token_limit
):
break
elif len(queried_former_failed_knowledge_to_render) > 1:
+19 -8
View File
@@ -4,11 +4,14 @@ from abc import ABC, abstractmethod
from collections.abc import Generator
from typing import TYPE_CHECKING, Any, Generic, TypeVar
from filelock import FileLock
from tqdm import tqdm
if TYPE_CHECKING:
from rdagent.core.evolving_framework import EvolvableSubjects
from contextlib import nullcontext
from rdagent.core.evaluation import EvaluableObj, Evaluator, Feedback
from rdagent.core.evolving_framework import EvolvingStrategy, EvoStep
from rdagent.log import rdagent_logger as logger
@@ -55,6 +58,8 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
with_knowledge: bool = False,
with_feedback: bool = True,
knowledge_self_gen: bool = False,
enable_filelock: bool = False,
filelock_path: str | None = None,
) -> None:
super().__init__(max_loop, evolving_strategy)
self.rag = rag
@@ -62,6 +67,8 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
self.with_knowledge = with_knowledge
self.with_feedback = with_feedback
self.knowledge_self_gen = knowledge_self_gen
self.enable_filelock = enable_filelock
self.filelock_path = filelock_path
def multistep_evolve(
self,
@@ -70,35 +77,39 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
) -> Generator[EvolvableSubjects, None, None]:
for evo_loop_id in tqdm(range(self.max_loop), "Implementing"):
with logger.tag(f"evo_loop_{evo_loop_id}"):
# 1. knowledge self-evolving
if self.knowledge_self_gen and self.rag is not None:
self.rag.generate_knowledge(self.evolving_trace)
# 2. RAG
# 1. RAG
queried_knowledge = None
if self.with_knowledge and self.rag is not None:
# TODO: Putting the evolving trace in here doesn't actually work
queried_knowledge = self.rag.query(evo, self.evolving_trace)
# 3. evolve
# 2. evolve
evo = self.evolving_strategy.evolve(
evo=evo,
evolving_trace=self.evolving_trace,
queried_knowledge=queried_knowledge,
)
# 4. Pack evolve results
# 3. Pack evolve results
es = EvoStep(evo, queried_knowledge)
# 5. Evaluation
# 4. Evaluation
if self.with_feedback:
es.feedback = (
eva if isinstance(eva, Feedback) else eva.evaluate(evo, queried_knowledge=queried_knowledge)
)
logger.log_object(es.feedback, tag="evolving feedback")
# 6. update trace
# 5. update trace
self.evolving_trace.append(es)
# 6. knowledge self-evolving
if self.knowledge_self_gen and self.rag is not None:
with FileLock(self.filelock_path) if self.enable_filelock else nullcontext(): # type: ignore[arg-type]
self.rag.load_dumped_knowledge_base()
self.rag.generate_knowledge(self.evolving_trace)
self.rag.dump_knowledge_base()
yield evo # yield the control to caller for process control and logging.
# 7. check if all tasks are completed
+21 -2
View File
@@ -77,8 +77,16 @@ class EvolvingStrategy(ABC):
class RAGStrategy(ABC):
"""Retrieval Augmentation Generation Strategy"""
def __init__(self, knowledgebase: EvolvingKnowledgeBase) -> None:
self.knowledgebase: EvolvingKnowledgeBase = knowledgebase
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.knowledgebase: EvolvingKnowledgeBase = self.load_or_init_knowledge_base(*args, **kwargs)
@abstractmethod
def load_or_init_knowledge_base(
self,
*args: Any,
**kwargs: Any,
) -> EvolvingKnowledgeBase:
pass
@abstractmethod
def query(
@@ -102,3 +110,14 @@ class RAGStrategy(ABC):
RAGStrategy should maintain the new knowledge all by itself.
"""
@abstractmethod
def dump_knowledge_base(self, *args: Any, **kwargs: Any) -> None:
pass
@abstractmethod
def load_dumped_knowledge_base(self, *args: Any, **kwargs: Any) -> None:
"""This is to load the dumped knowledge base.
It's mainly used in parallel coding of which several coder shares the same knowledge base.
Then the agent should load the knowledge base from others before updating it.
"""
+4
View File
@@ -681,3 +681,7 @@ class APIBackend(ABC):
Call the chat completion function
"""
raise NotImplementedError("Subclasses must implement this method")
@property
def chat_token_limit(self) -> int:
return LLM_SETTINGS.chat_token_limit
+11
View File
@@ -7,6 +7,7 @@ from litellm import (
completion,
completion_cost,
embedding,
get_max_tokens,
supports_function_calling,
supports_response_schema,
token_counter,
@@ -201,3 +202,13 @@ class LiteLLMAPIBackend(APIBackend):
Check if the backend supports function calling
"""
return supports_response_schema(model=LITELLM_SETTINGS.chat_model) and LITELLM_SETTINGS.enable_response_schema
@property
def chat_token_limit(self) -> int:
try:
max_tokens = get_max_tokens(LITELLM_SETTINGS.chat_model)
if max_tokens is None:
return super().chat_token_limit
return max_tokens
except Exception as e:
return super().chat_token_limit
@@ -61,17 +61,17 @@ class DSRunnerMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
}
output_spec, extract_output_fn = output_map[self.settings.diff_mode]
if prev_task_feedback.hyperparameter_tuning_decision:
# Use system_refine for hyperparameter tuning
system_prompt = T(".prompts:DSCoSTEER.system_refine").r(
if prev_task_feedback.final_decision is False:
task_information_str = target_task.get_task_information()
# Use system_debugger for error fixing and debugging
system_prompt = T(".prompts:DSCoSTEER.system_debugger").r(
task_desc=task_information_str,
out_spec=output_spec,
diff_mode=self.settings.diff_mode,
)
else:
task_information_str = target_task.get_task_information()
# Use system_debugger for error fixing and debugging
# Use system_refine for hyperparameter tuning
system_prompt = T(".prompts:DSCoSTEER.system_refine").r(
task_desc=task_information_str,
out_spec=output_spec,
diff_mode=self.settings.diff_mode,
)
@@ -152,12 +152,13 @@ class DSCoSTEERRunner(CoSTEER):
)
def develop(self, exp):
bak_sub_tasks = exp.sub_tasks
bak_sub_tasks = exp.pending_tasks_list
exp.sub_tasks = [
CoSTEERTask(
name="Debug running solution",
description=f"You'll be provided with the source code and the running and testing stdout. "
"Please check the error messages and debug the source code if any errors occur.\n"
f"Original task: {bak_sub_tasks[0][0].get_task_information()}\n"
f"Current code repo md5: {md5_hash(exp.experiment_workspace.all_codes)}",
),
]
@@ -50,6 +50,10 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator):
queried_knowledge: QueriedKnowledge = None,
**kwargs,
) -> DSCoSTEEREvalFeedback:
if len(queried_knowledge.task_to_former_failed_traces[target_task.get_task_information()][0]) == 0:
enable_hyperparameter_tuning_check = True
else:
enable_hyperparameter_tuning_check = False
env = get_ds_env(
extra_volumes={
@@ -133,6 +137,7 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator):
scenario=self.scen.get_scenario_all_desc(eda_output=implementation.file_dict.get("EDA.md", None)),
is_sub_enabled=test_eval.is_sub_enabled(self.scen.competition),
task_desc=target_task.get_task_information(),
enable_hyperparameter_tuning_check=enable_hyperparameter_tuning_check,
)
user_prompt = T(".prompts:DSCoSTEER_eval.user").r(
code=implementation.all_codes,
@@ -25,7 +25,7 @@ DSCoSTEER_eval:
3. Confirm that the prediction file (`submission.csv`) is generated using only the test dataset, and its format matches the sample submission.
If the code does not satisfy the requirements:
- Set "final_decision" to false.
- set "hyperparameter_tuning_decision" to false.
{% if enable_hyperparameter_tuning_check %}- set "hyperparameter_tuning_decision" to false.
- Set "hyperparameter_tuning_suggestion" to an empty string.
If the code satisfy the requirements:
- Proceed to the next evaluation.
@@ -48,6 +48,7 @@ DSCoSTEER_eval:
If the code does not satisfy the requirements:
- Set "hyperparameter_tuning_decision" to false.
- Set "hyperparameter_tuning_suggestion" to an empty string.
{% endif %}
## Output format
Please respond with your feedback in the following JSON format and order without anything else:
@@ -55,15 +56,36 @@ DSCoSTEER_eval:
{
"execution": "Describe whether the whole code base executed successfully and generating the final submission. Include any errors or issues encountered, and retain all error messages and traceback details.",
"return_checking": "Verify the generated files, particularly the submission file. Ensure that its format matches the sample submission",
"code": "Provide feedback on code quality, readability, and adherence to the given specifications.",
"code": "Provide feedback on code quality, readability, and adherence to the given specifications.",{% if enable_hyperparameter_tuning_check %}
"hyperparameter_tuning_decision": <true/false>,
"hyperparameter_tuning_suggestion": <suggestion in plain text for hyperparameter tuning>,
"hyperparameter_tuning_suggestion": <suggestion in plain text for hyperparameter tuning>,{% endif %}
"final_decision": <true/false>,
}
```
{% else %}
The user will provide you the whole code base, some logs generated during the execution of the whole workflow. Your evaluation scope includes whether the workflow code executes successfully.
No need to check the detail of submission file.
{% if enable_hyperparameter_tuning_check %}
# Evaluation: Hyperparameter
## Evaluation Description
The user will provide you the time spent on the whole code execution and the timeout of the code execution. You should decide whether the hyperparameter is reasonable based on the time.
If the code uses only a very small portion (below 25%) of the allowed time, and hyperparameters like `n_estimators` or `epochs` have low values, with early stopping not being triggered and possible signs of underfitting, you should suggest increasing these hyperparameters.
You should also notice other resources utilization hyper-parameters.
For example, if you are using a GPU with large memory, and the batch size is set very low, you should suggest increasing the batch size if it is not reasonable.
## Evaluation Guidelines
1. The code execution time or resource utilization suggest that there is room for improvement in the hyperparameters.
2. The code must apply early stopping strategy already (in order to prevent overfitting).
3. Your suggestion should have a strong chance of improving the model's performance. Focus on the most obvious and impactful opportunities for quick improvement by leveraging more training time. Don't explore hyperparameters with low confidence. If there are no obvious and impactful opportunities and the code runs well, please accept it.
If the code satisfy the requirements:
- Set "hyperparameter_tuning_decision" to true.
- Set "final_decision" to false.
- Provide a reasonable suggestion in "hyperparameter_tuning_suggestion". The "hyperparameter_tuning_suggestion" should begin with a clear observation, followed by your suggestion. For example: "[Observation] The maximum number of epochs was reached, but the validation loss is still going down and early stopping was not activated. Only 15% of the allowed time was used. [Suggestion] We recommend increasing epochs to 100 to avoid underfitting and further improve model performance."
If the code does not satisfy the requirements:
- Set "hyperparameter_tuning_decision" to false.
- Set "hyperparameter_tuning_suggestion" to an empty string.
{% endif %}
Please respond with your feedback in the following JSON format and order
```json
@@ -71,8 +93,8 @@ DSCoSTEER_eval:
"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.",
"return_checking": "Describe the expected file to be generated.",
"code": "Provide feedback on code quality, readability, and adherence to the given specifications.",
"hyperparameter_tuning_decision": <true/false>,
"hyperparameter_tuning_suggestion": <suggestion in plain text for hyperparameter tuning>,
{% if enable_hyperparameter_tuning_check %}"hyperparameter_tuning_decision": <true/false>,
"hyperparameter_tuning_suggestion": <suggestion in plain text for hyperparameter tuning>,{% endif %}
"final_decision": <true/false>,
}
```
+15 -7
View File
@@ -31,6 +31,7 @@ from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback
from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
from rdagent.scenarios.data_science.proposal.exp_gen import DSTrace
from rdagent.scenarios.data_science.proposal.exp_gen.base import DataScienceScen
from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase
from rdagent.scenarios.data_science.proposal.exp_gen.proposal import DSProposalV2ExpGen
from rdagent.utils.workflow.misc import wait_retry
@@ -39,17 +40,21 @@ from rdagent.utils.workflow.misc import wait_retry
def clean_workspace(workspace_root: Path) -> None:
"""
Clean the workspace folder and only keep the essential files to save more space.
workspace_root might contain a file in parallel with the folders, we should directly remove it.
# remove all files and folders in the workspace except for .py, .md, and .csv files to avoid large workspace dump
"""
for file_and_folder in workspace_root.iterdir():
if file_and_folder.is_dir():
if file_and_folder.is_symlink():
if workspace_root.is_file():
workspace_root.unlink()
else:
for file_and_folder in workspace_root.iterdir():
if file_and_folder.is_dir():
if file_and_folder.is_symlink():
file_and_folder.unlink()
else:
shutil.rmtree(file_and_folder)
elif file_and_folder.is_file() and file_and_folder.suffix not in [".py", ".md", ".csv"]:
file_and_folder.unlink()
else:
shutil.rmtree(file_and_folder)
elif file_and_folder.is_file() and file_and_folder.suffix not in [".py", ".md", ".csv"]:
file_and_folder.unlink()
@wait_retry()
@@ -215,6 +220,9 @@ class DataScienceRDLoop(RDLoop):
self.trace.sync_dag_parent_and_hist((exp, prev_out["feedback"]), cur_loop_id)
else:
exp: DSExperiment = prev_out["direct_exp_gen"] if isinstance(e, CoderError) else prev_out["coding"]
# TODO: distinguish timeout error & other exception.
if isinstance(self.trace.scen, DataScienceScen) and DS_RD_SETTING.allow_longer_timeout:
self.trace.scen.increase_timeout()
# set the local selection to the trace as global selection, then set the DAG parent for the trace
if exp.local_selection is not None:
@@ -22,6 +22,7 @@ class DSHypothesis(Hypothesis):
problem_name: str | None = None,
problem_desc: str | None = None,
problem_label: Literal["SCENARIO_PROBLEM", "FEEDBACK_PROBLEM"] = "FEEDBACK_PROBLEM",
appendix: str | None = None,
) -> None:
super().__init__(
hypothesis, reason, concise_reason, concise_observation, concise_justification, concise_knowledge
@@ -30,6 +31,7 @@ class DSHypothesis(Hypothesis):
self.problem_name = problem_name
self.problem_desc = problem_desc
self.problem_label = problem_label
self.appendix = appendix
def __str__(self) -> str:
if self.hypothesis is None:
@@ -44,6 +46,8 @@ class DSHypothesis(Hypothesis):
lines.append(f"Hypothesis: {self.hypothesis}")
if self.reason is not None:
lines.append(f"Reason: {self.reason}")
if self.appendix is not None:
lines.append(f"Appendix: {self.appendix}")
return "\n".join(lines)
@@ -306,6 +306,18 @@ hypothesis_rewrite:
3. **Decisive Specificity** Remove all ambiguity; propose one clear action.
4. **Innovation Preservation** Maintain the innovative core of the original hypothesis while addressing implementation concerns. Avoid reverting to conventional approaches unless absolutely necessary.
5. **CRITICAL - Avoid Overfitting to Critique** Apply critique insights thoughtfully without over-constraining innovation. Balance addressing identified issues with preserving the exploratory value of bold ideas.
{% if enable_scale_check %}6. The user is currently working on a continuous exploration on the task. It's typical that we first try in small scale and in some certain point we will scale up the solution.
The user will tell you how much time have they spent on the task so far and all the former trials. You should consider whether to scale up the solution based on the current situation. You should put this conclusion in each hypothesis's appendix section.
Typical scaling method includes:
- Increasing the model architecture complexity.
- Increasing the number of models to ensemble.
- Increasing the number of features.
- Increasing the number of cross validation folds.
- Increasing the number of epochs for training.
- Increasing the batch size for training.
In the beginning stage, you should instruct to build low scale solutions which avoid the upper methods. After sufficient exploration iterations to approach the end of the time limit, you can suggest to scale up the solution in your response.
Scaling is no connection to the debugging process. It's related to the whole solution's complexity. Please include this in every hypothesis you rewrite.
{% endif %}
## Guidelines for Writing Rewritten Hypotheses
@@ -362,6 +374,11 @@ hypothesis_rewrite:
# Original Hypotheses and Their Critiques
{{ hypothesis_critique_pairs }}
{% if time_status is not none %}
# Time Status
{{ time_status }}
{% endif %}
task_gen:
system: |-
@@ -400,6 +417,8 @@ task_gen:
- Handle missing values and outliers appropriately as guided by the hypothesis or SOTA.
- Ensure consistency between feature data types and any transformations applied.
- Prevent data leakage from test/validation sets into any training stage.
- Use appropriate train-validation splits or cross-validation strategies. Some dataset might not be suitable for StratifiedShuffleSplit since some categories may not be present in the test set. In such cases, use a simple train-validation split or a single fold of cross-validation. Implement a try except block to handle potential errors if you are using StratifiedShuffleSplit.
- Use appropriate cross-validation strategies. Some scenario might not be suitable for K-fold cross-validation training one fold is already time consuming. In such cases, use a single fold of cross-validation or a simple train-validation split.
6. **Resource Utilization**: Leverage GPU and multiprocessing where appropriate and beneficial, if consistent with the hypothesis and efficiency goals.
7. **Metric Calculation and Storage (`scores.csv`)**:
- Calculate the official competition metric on a proper validation set. Save results to `scores.csv`.
@@ -629,6 +648,7 @@ output_format:
"reason": "Independent justification for why this hypothesis makes sense given the current scenario, dataset characteristics, and competition requirements. DO NOT reference critique feedback or suggestions. Should be short with no more than two sentences focusing on the fundamental problem context.",
"component": "The component tag of the hypothesis. Must be one of ('DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow').",
"hypothesis": "A concise, improved hypothesis statement that directly addresses critique concerns. Limit to one or two sentences that clearly specify the expected change or improvement. Should be more specific and actionable than the original.",
{% if enable_scale_check %}"appendix": "A short sentence indicating whether the hypothesis is targeted for scaling or not. Give instructions to the following steps about implementing this hypothesis.", {% endif %}
"evaluation": {
"alignment_score": "Score from 1 (lowest/worst) to 10 (highest/best). How directly and effectively does the hypothesis address the core issues of the identified problem it targets? A higher score means a stronger, more direct alignment.",
"impact_score": "Score from 1 (lowest/worst) to 10 (highest/best). What is the estimated magnitude of improvement (e.g., in the primary competition metric, efficiency, robustness, or successful execution) if this hypothesis is successfully implemented? Higher scores for greater positive impact.",
@@ -23,7 +23,10 @@ from rdagent.scenarios.data_science.proposal.exp_gen.draft.draft import (
DSDraftExpGen, # TODO: DSDraftExpGen should be moved to router in the further
)
from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSIdea
from rdagent.scenarios.data_science.proposal.exp_gen.planner import DSExperimentPlan
from rdagent.scenarios.data_science.proposal.exp_gen.planner import (
DSExperimentPlan,
RD_Agent_TIMER_wrapper,
)
from rdagent.scenarios.data_science.proposal.exp_gen.utils import get_packages
from rdagent.utils.agent.tpl import T
from rdagent.utils.repo.diff import generate_diff_from_dict
@@ -742,14 +745,28 @@ class DSProposalV2ExpGen(ExpGen):
hypothesis_critique_pairs += f"**Reasoning:** {hypothesis_data.get('reason', 'Not provided')}\n"
hypothesis_critique_pairs += f"**Critique:** {critique_data.get('critique', 'No critique available')}\n\n"
time_status = None
if DS_RD_SETTING.enable_scale_check and RD_Agent_TIMER_wrapper.timer.started:
remain_time = RD_Agent_TIMER_wrapper.timer.remain_time()
all_duration = RD_Agent_TIMER_wrapper.timer.all_duration
remain_percent = remain_time / all_duration
time_status = (
f"Remain time: {remain_time.total_seconds() / 3600:.2f} hours, "
f"{remain_percent:.2%} remaining of total time: {all_duration.total_seconds() / 3600:.2f} hours."
)
sys_prompt = T(".prompts_v2:hypothesis_rewrite.system").r(
rewrite_output_format=T(".prompts_v2:output_format.rewrite").r(),
rewrite_output_format=T(".prompts_v2:output_format.rewrite").r(
enable_scale_check=DS_RD_SETTING.enable_scale_check
),
enable_scale_check=DS_RD_SETTING.enable_scale_check,
)
user_prompt = T(".prompts_v2:hypothesis_rewrite.user").r(
scenario_desc=scenario_desc,
exp_and_feedback_list_desc=exp_feedback_list_desc,
sota_exp_desc=sota_exp_desc,
hypothesis_critique_pairs=hypothesis_critique_pairs,
time_status=time_status,
)
response = APIBackend().build_messages_and_create_chat_completion(
@@ -763,7 +780,12 @@ class DSProposalV2ExpGen(ExpGen):
# Validate that we have rewritten hypotheses for all original hypotheses
expected_problems = set(hypothesis_dict.keys())
available_problems = set(improved_hypotheses_dict.keys())
available_problems = set( # The code snippet provided is a comment in Python. It appears to be
# a placeholder for a function or variable named
# `improved_hypotheses_dict`. The actual implementation of this
# function or variable is not provided in the code snippet.
improved_hypotheses_dict.keys()
)
if not expected_problems.issubset(available_problems):
missing_problems = expected_problems - available_problems
@@ -867,6 +889,7 @@ class DSProposalV2ExpGen(ExpGen):
problem_name=max_score_problem_name,
problem_desc=problem_dict.get("problem", "Problem description not provided"),
problem_label=problem_dict.get("label", "FEEDBACK_PROBLEM"),
appendix=hypothesis_dict[max_score_problem_name].get("appendix", None),
)
def task_gen(
@@ -965,6 +988,7 @@ class DSProposalV2ExpGen(ExpGen):
problem_name=name,
problem_desc=problem_data.get("problem", "Problem description not provided"),
problem_label=problem_data.get("label", "FEEDBACK_PROBLEM"),
appendix=data.get("appendix", None),
)
)
return result
@@ -114,21 +114,34 @@ class DataScienceScen(Scenario):
self.longer_time_limit_required = response_json_analysis.get(
"Longer time limit required", False
) # True or False, whether the competition scenario requires a longer time limit to the code.
self.timeout_increase_count = 0
def real_debug_timeout(self):
return (
DS_RD_SETTING.debug_timeout * DS_RD_SETTING.coder_longer_timeout_multiplier
DS_RD_SETTING.debug_timeout
* min(
DS_RD_SETTING.coder_longer_timeout_multiplier_upper,
self.timeout_increase_count * DS_RD_SETTING.timeout_increase_stage + 1,
)
if self.longer_time_limit_required and DS_RD_SETTING.allow_longer_timeout
else DS_RD_SETTING.debug_timeout
)
def real_full_timeout(self):
return (
DS_RD_SETTING.full_timeout * DS_RD_SETTING.runner_longer_timeout_multiplier
DS_RD_SETTING.full_timeout
* min(
DS_RD_SETTING.runner_longer_timeout_multiplier_upper,
self.timeout_increase_count * DS_RD_SETTING.timeout_increase_stage + 1,
)
if self.longer_time_limit_required and DS_RD_SETTING.allow_longer_timeout
else DS_RD_SETTING.full_timeout
)
def increase_timeout(self):
"""Increase the timeout multiplier for the scenario."""
self.timeout_increase_count += 1
@property
def background(self) -> str:
background_template = T(".prompts:competition_background")
@@ -82,9 +82,9 @@ def classify_report_from_dict(
user_prompt=content,
system_prompt=classify_prompt,
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
content = content[: -(LLM_SETTINGS.chat_token_limit // 100)]
content = content[: -(APIBackend().chat_token_limit // 100)]
vote_list = []
for _ in range(vote_time):
@@ -367,7 +367,7 @@ def __check_factor_duplication_simulate_json_mode(
APIBackend().build_messages_and_calculate_token(
user_prompt=current_df.to_string(), system_prompt=T(".prompts:factor_duplicate_system").r()
)
> LLM_SETTINGS.chat_token_limit
> APIBackend().chat_token_limit
):
working_list.append(current_df.iloc[: current_df.shape[0] // 2, :])
working_list.append(current_df.iloc[current_df.shape[0] // 2 :, :])
+13 -4
View File
@@ -121,6 +121,15 @@ def filter_redundant_text(stdout: str) -> str:
# Collapse any excessive blank lines/spaces
filtered_stdout = try_regex_sub(r"\s*\n\s*", filtered_stdout, replace_with="\n")
# remove repeated lines
lines_to_count: dict[str, int] = {}
filtered_stdout_lines = filtered_stdout.splitlines()
for line in filtered_stdout_lines:
lines_to_count[line] = lines_to_count.get(line, 0) + 1
filtered_stdout = "\n".join(
[line for line in filtered_stdout_lines if lines_to_count[line] <= len(filtered_stdout_lines) // 10]
)
# Iteratively ask the LLM for additional filtering patterns (up to 3 rounds)
for _ in range(3):
truncated_stdout = filtered_stdout
@@ -133,11 +142,11 @@ def filter_redundant_text(stdout: str) -> str:
user_prompt=user_prompt,
system_prompt=system_prompt,
)
if stdout_token_size < LLM_SETTINGS.chat_token_limit * 0.1:
if stdout_token_size < APIBackend().chat_token_limit * 0.1:
return truncated_stdout
elif stdout_token_size > LLM_SETTINGS.chat_token_limit * 0.6:
head = truncated_stdout[: int(LLM_SETTINGS.chat_token_limit * 0.3)]
tail = truncated_stdout[-int(LLM_SETTINGS.chat_token_limit * 0.3) :]
elif stdout_token_size > APIBackend().chat_token_limit * 0.6:
head = truncated_stdout[: int(APIBackend().chat_token_limit * 0.3)]
tail = truncated_stdout[-int(APIBackend().chat_token_limit * 0.3) :]
truncated_stdout = head + tail
else:
break