Ability to override clean_code_for_run (#28783)

* Add clean_code_for_run function

* Call clean_code_for_run from agent method
This commit is contained in:
w4ffl35 2024-02-05 02:48:41 +00:00 committed by GitHub
parent c430d6eaee
commit 0466fd5ca2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 3 deletions

View File

@ -315,6 +315,13 @@ class Agent:
self.chat_state = {}
self.cached_tools = None
def clean_code_for_run(self, result):
"""
Override this method if you want to change the way the code is
cleaned for the `run` method.
"""
return clean_code_for_run(result)
def run(self, task, *, return_code=False, remote=False, **kwargs):
"""
Sends a request to the agent.
@ -339,7 +346,7 @@ class Agent:
"""
prompt = self.format_prompt(task)
result = self.generate_one(prompt, stop=["Task:"])
explanation, code = clean_code_for_run(result)
explanation, code = self.clean_code_for_run(result)
self.log(f"==Explanation from the agent==\n{explanation}")

View File

@ -14,7 +14,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .agents import BASE_PYTHON_TOOLS, clean_code_for_chat, clean_code_for_run
from .agents import BASE_PYTHON_TOOLS, clean_code_for_chat
from .python_interpreter import InterpretorError, evaluate
@ -554,7 +554,7 @@ def evaluate_agent(agent, batch_size=8, verbose=False, return_errors=False):
problem = EVALUATION_TASKS[eval_idx[start_idx + idx]]
if verbose:
print(f"====Task {start_idx + idx}====\n{batch_tasks[idx]}\n")
explanation, code = clean_code_for_run(result)
explanation, code = agent.clean_code_for_run(result)
# Evaluate agent answer and code answer
agent_answer = evaluate_code(code, problem.inputs, verbose=verbose)