Files
NexQuant/rdagent/utils/agent/ret.py
T
XianBW e21b334741 feat: out spec change for o1-preview (#666)
* add not_json batcheditout

* ensemble out_spec change

* feature out_spec change

* model out_spec change

* workflow out_spec change

* runner debugger out_spec change

* filter_progress_bar return format fix

* data_loader and spec out_spec change

* show finish_reason in llm log

* json_mode fix

* remove hardcode

* fix CI

* fix grammer

* complete PythonBatchEditOut logic

---------

Co-authored-by: yuanteli <1957922024@qq.com>
2025-03-18 17:43:44 +08:00

70 lines
1.6 KiB
Python

"""
The output of a agent is very important.
We think this part can be shared.
"""
import json
import re
from abc import abstractclassmethod
from typing import Any
from rdagent.utils.agent.tpl import T
class AgentOut:
json_mode: bool = False # To get the output, is json_mode required.
@abstractclassmethod
def get_spec(cls, **context: Any) -> str:
raise NotImplementedError(f"Please implement the `get_spec` method")
@classmethod
def extract_output(cls, resp: str) -> Any:
raise resp
class PythonAgentOut(AgentOut):
@classmethod
def get_spec(cls):
return T(".tpl:PythonAgentOut").r()
@classmethod
def extract_output(cls, resp: str):
match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL)
if match:
code = match.group(1)
return code
return resp
class BatchEditOut(AgentOut):
json_mode: bool = True
@classmethod
def get_spec(cls, with_del=True):
return T(".tpl:BatchEditOut").r(with_del=with_del)
@classmethod
def extract_output(cls, resp: str):
return json.loads(resp)
class PythonBatchEditOut(AgentOut):
@classmethod
def get_spec(cls, with_del=True):
return T(".tpl:PythonBatchEditOut").r(with_del=with_del)
@classmethod
def extract_output(cls, resp: str):
code_blocks = {}
pattern = re.compile(r"```(.*?)\n(.*?)\n```", re.DOTALL)
matches = pattern.findall(resp)
for match in matches:
file_name, code = match
code_blocks[file_name.strip()] = code.strip()
return code_blocks