初始化项目，由ModelHub XC社区提供模型
Model: bk1dr/qwen3-8b-code-pkpo Source: Original Platform
This commit is contained in:
ModelHub XC
273
agent_core.py
Normal file
273
agent_core.py
Normal file
@@ -0,0 +1,273 @@
|
||||
"""Shared agent core: chat template, instructions, tool-call parsing, episode logic.
|
||||
|
||||
This exact module drives (1) RL training rollouts, (2) every evaluation, and
|
||||
(3) the shipped standalone tool — zero train/inference mismatch.
|
||||
|
||||
Conventions (inherited from the reference client):
|
||||
- NO system role: instructions are merged into the first user message.
|
||||
- Strict user/assistant alternation; tool output comes back as the next user
|
||||
message, prefixed "Tool output: ".
|
||||
- Plain-text tool calls, types "bash" and "final", lenient regex parsing, and
|
||||
parse errors are fed back to the model as the next user message.
|
||||
|
||||
Format: reasoning goes in <think>...</think>, the tool call in <answer>...</answer>
|
||||
(the reference's "Thoughts:" field is dropped — it is redundant with <think>).
|
||||
Assistant turns end with "</answer>" followed by EOS. Past assistant turns are
|
||||
rendered with the <think> block stripped (saves context; template handles it).
|
||||
"""
|
||||
import re
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Chat template
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
PREAMBLE = (
|
||||
"A conversation between User and Assistant. The user asks a question, and the "
|
||||
"Assistant solves it. The assistant first thinks about the reasoning process in "
|
||||
"the mind and then provides the user with the answer. The reasoning process and "
|
||||
"answer are enclosed within <think>...</think> and <answer>...</answer> tags, "
|
||||
"respectively, i.e., <think> reasoning process here </think> "
|
||||
"<answer>answer here </answer>."
|
||||
)
|
||||
|
||||
# Past assistant messages are rendered with everything up to and including
|
||||
# </think> removed. The generation prompt ends with "Assistant: <think>" + newline.
|
||||
CHAT_TEMPLATE = (
|
||||
"{{- '" + PREAMBLE.replace("'", "\\'") + "\\n\\n' -}}"
|
||||
"{%- for message in messages -%}"
|
||||
"{%- if message['role'] == 'user' -%}"
|
||||
"{{- 'User: ' + message['content'] + '\\n\\n' -}}"
|
||||
"{%- elif message['role'] == 'assistant' -%}"
|
||||
"{%- set content = message['content'] -%}"
|
||||
"{%- if '</think>' in content -%}"
|
||||
"{%- set content = content.split('</think>')[-1] -%}"
|
||||
"{%- endif -%}"
|
||||
"{{- 'Assistant: ' + content.strip() + '\\n\\n' -}}"
|
||||
"{%- endif -%}"
|
||||
"{%- endfor -%}"
|
||||
"{%- if add_generation_prompt -%}"
|
||||
"{{- 'Assistant: <think>\\n' -}}"
|
||||
"{%- endif -%}"
|
||||
)
|
||||
|
||||
THINK_PREFIX = "<think>\n" # generation prompt ends with this; completions start after it
|
||||
STOP_STRING = "</answer>"
|
||||
|
||||
# Sampling defaults — used for RL rollouts, every eval, and shipped generation_config.
|
||||
GEN_TEMPERATURE = 1.0
|
||||
GEN_TOP_P = 0.95
|
||||
MAX_TURN_TOKENS = 1024 # shared training/eval/shipped-tool per-turn output budget
|
||||
|
||||
|
||||
def render_messages(messages, add_generation_prompt: bool = True) -> str:
|
||||
"""Render messages exactly like CHAT_TEMPLATE, without depending on Jinja.
|
||||
|
||||
Offline vLLM/Transformers calls use this function; server-mode vLLM uses the
|
||||
tokenizer chat template below. Keep both byte-identical.
|
||||
"""
|
||||
parts = [PREAMBLE + "\n\n"]
|
||||
for message in messages:
|
||||
role = message["role"]
|
||||
content = message["content"]
|
||||
if role == "user":
|
||||
parts.append("User: " + content + "\n\n")
|
||||
elif role == "assistant":
|
||||
if "</think>" in content:
|
||||
content = content.split("</think>")[-1]
|
||||
parts.append("Assistant: " + content.strip() + "\n\n")
|
||||
if add_generation_prompt:
|
||||
parts.append("Assistant: <think>\n")
|
||||
return "".join(parts)
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Instructions (merged into the first user message)
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
TOOL_INSTRUCTIONS = """You are an AI agentic coding assistant. You complete the user's task by calling tools step by step, one tool call per turn.
|
||||
|
||||
First reason inside <think>...</think>, then put exactly one tool call inside <answer>...</answer> in exactly this format:
|
||||
|
||||
<answer>Tool type: [bash or final]
|
||||
Tool query: [the shell command to run, or the final answer]</answer>
|
||||
|
||||
Available tools:
|
||||
1. bash: runs a shell command in an isolated Linux sandbox and returns its output (e.g. write files with heredocs, run python, test code). The sandbox is stateless: the working directory and variables reset between calls, so chain commands with && where needed.
|
||||
2. final: ends the task; put the complete final answer in Tool query.
|
||||
|
||||
The next user message will contain the tool output as "Tool output: ...". Keep outputs small (use head/tail). After your final tool call the episode ends."""
|
||||
|
||||
|
||||
def coding_task_message(statement: str, starter_code: str = "", max_turns: int = 2) -> str:
|
||||
"""First user message for a competitive-programming task (training AND eval)."""
|
||||
if starter_code and starter_code.strip():
|
||||
req = (
|
||||
"Write a Python 3 solution that completes the following starter code, "
|
||||
"keeping the exact same class name and method signature:\n"
|
||||
"```python\n" + starter_code.strip() + "\n```"
|
||||
)
|
||||
else:
|
||||
req = ("Write a complete Python 3 program that reads the input from stdin "
|
||||
"and prints the required output to stdout.")
|
||||
n_bash = max_turns - 1
|
||||
one_turn_rule = (
|
||||
"- This is a one-turn episode: NEVER call `bash`; submit the complete Python solution "
|
||||
"with your single `final` tool call immediately.\n"
|
||||
if max_turns == 1 else ""
|
||||
)
|
||||
return f"""{TOOL_INSTRUCTIONS}
|
||||
|
||||
Task: solve this competitive programming problem in Python 3.
|
||||
|
||||
{statement.strip()}
|
||||
|
||||
{req}
|
||||
|
||||
Rules:
|
||||
- You have at most {max_turns} tool calls in total, so at most {n_bash} bash call{"s" if n_bash != 1 else ""}; you may use bash to test your solution on the example input before submitting.
|
||||
{one_turn_rule}- Your last tool call must be "final" and its Tool query must contain ONLY your complete solution inside one ```python ... ``` code block.
|
||||
- The solution is judged against hidden tests; handle all edge cases and stay within a few seconds of runtime."""
|
||||
|
||||
|
||||
PARSE_ERROR_FEEDBACK = """Parsing Error: {err}
|
||||
Please repeat, strictly matching this format:
|
||||
|
||||
<answer>Tool type: [bash or final]
|
||||
Tool query: [the shell command, or the final answer]</answer>"""
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Parsing
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_action(text: str):
|
||||
"""Parse an assistant completion into (tool_type, tool_query, error).
|
||||
|
||||
`text` is the completion text (reasoning + <answer> block, with or without
|
||||
the closing tag). Lenient: falls back to scanning the whole text.
|
||||
"""
|
||||
m = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
|
||||
if m:
|
||||
seg = m.group(1)
|
||||
elif "<answer>" in text:
|
||||
seg = text.split("<answer>", 1)[1]
|
||||
else:
|
||||
seg = text.split("</think>")[-1]
|
||||
|
||||
tt = re.search(r"Tool type:\s*(.*?)(?=\n\s*Tool query:|$)", seg, re.DOTALL | re.IGNORECASE)
|
||||
tq = re.search(r"Tool query:\s*(.*)", seg, re.DOTALL | re.IGNORECASE)
|
||||
tool_type = tt.group(1).strip().lower() if tt else None
|
||||
# The reference protocol presents its choices as ``[bash or final]`` and
|
||||
# base completions commonly echo the brackets (``Tool type: [final]``).
|
||||
# Treat that as the same lenient plain-text call rather than wasting a turn.
|
||||
if tool_type:
|
||||
tool_type = tool_type.strip(" \t[](){}<>`*_\"'")
|
||||
tool_query = tq.group(1).strip() if tq else None
|
||||
|
||||
if not tool_type or tool_query is None or tool_query == "":
|
||||
return None, None, ("Could not parse 'Tool type' or 'Tool query' from the response. "
|
||||
"Ensure you output an <answer> block with 'Tool type:' and 'Tool query:'.")
|
||||
if tool_type not in ("bash", "final"):
|
||||
return None, None, f"Unknown tool type '{tool_type}'; must be 'bash' or 'final'."
|
||||
if tool_type == "bash":
|
||||
# strip a markdown fence around the command, reference-style
|
||||
tool_query = re.sub(r"^```[a-zA-Z]*\n", "", tool_query)
|
||||
tool_query = re.sub(r"\n```$", "", tool_query)
|
||||
tool_query = tool_query.strip()
|
||||
return tool_type, tool_query, None
|
||||
|
||||
|
||||
def extract_code(final_answer: str) -> str:
|
||||
"""Extract the python solution from a final answer (last ```python block, else raw)."""
|
||||
blocks = re.findall(r"```(?:python|py)?\n(.*?)```", final_answer, re.DOTALL)
|
||||
if blocks:
|
||||
return blocks[-1].strip()
|
||||
return final_answer.strip()
|
||||
|
||||
|
||||
def normalize_completion(content: str, finish_reason: str = "stop") -> str:
|
||||
"""Make server-mode completions byte-identical to training rollouts: the stop
|
||||
string is part of the turn. If the server stripped it (or stopped at EOS right
|
||||
after it), re-append."""
|
||||
if STOP_STRING not in content and finish_reason == "stop" and "<answer>" in content:
|
||||
content = content.rstrip() + STOP_STRING
|
||||
return content
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Episode state machine (transport-agnostic)
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Episode:
|
||||
"""One agentic episode with strict user/assistant alternation.
|
||||
|
||||
Use: messages -> (render+generate outside) -> step(completion) which returns
|
||||
("bash", cmd) | ("final", answer) | ("continue", None) after appending messages.
|
||||
After executing a bash command, call add_tool_output(out).
|
||||
"""
|
||||
|
||||
def __init__(self, first_user_msg: str, max_turns: int = 2, meta=None):
|
||||
self.messages = [{"role": "user", "content": first_user_msg}]
|
||||
self.max_turns = max_turns
|
||||
self.turns = 0
|
||||
self.done = False
|
||||
self.final_answer = None
|
||||
self.end_reason = None
|
||||
self.meta = meta or {}
|
||||
|
||||
def step(self, completion: str):
|
||||
"""completion: model output text (starts after 'Assistant: <think>\\n')."""
|
||||
assert not self.done
|
||||
content = THINK_PREFIX + completion
|
||||
self.messages.append({"role": "assistant", "content": content})
|
||||
self.turns += 1
|
||||
tool_type, tool_query, err = parse_action(completion)
|
||||
|
||||
if err is not None:
|
||||
if self.turns >= self.max_turns:
|
||||
self.done, self.end_reason = True, "parse_error"
|
||||
return ("continue", None)
|
||||
self.add_tool_output(PARSE_ERROR_FEEDBACK.format(err=err))
|
||||
return ("continue", None)
|
||||
|
||||
if tool_type == "final":
|
||||
self.done, self.end_reason = True, "final"
|
||||
self.final_answer = tool_query
|
||||
return ("final", tool_query)
|
||||
|
||||
# bash
|
||||
if self.turns >= self.max_turns:
|
||||
self.done, self.end_reason = True, "out_of_turns"
|
||||
return ("continue", None)
|
||||
return ("bash", tool_query)
|
||||
|
||||
def add_tool_output(self, output: str):
|
||||
self.messages.append({"role": "user", "content": f"Tool output: {output}"})
|
||||
|
||||
|
||||
def run_episode(generate_fn, bash_fn, first_user_msg: str, max_turns: int = 2, meta=None):
|
||||
"""Sequential episode driver (eval + shipped tool). generate_fn(messages)->completion."""
|
||||
ep = Episode(first_user_msg, max_turns=max_turns, meta=meta)
|
||||
while not ep.done:
|
||||
completion = generate_fn(ep.messages)
|
||||
kind, arg = ep.step(completion)
|
||||
if kind == "bash":
|
||||
ep.add_tool_output(bash_fn(arg))
|
||||
return ep
|
||||
|
||||
|
||||
def openai_generate_fn(client, model: str, temperature: float = GEN_TEMPERATURE,
|
||||
top_p: float = GEN_TOP_P, max_tokens: int = MAX_TURN_TOKENS):
|
||||
"""generate_fn over an OpenAI-compatible chat endpoint (vLLM server). The server
|
||||
applies this repo's chat template (shipped in the tokenizer)."""
|
||||
def fn(messages):
|
||||
kwargs = dict(model=model, messages=messages, temperature=temperature,
|
||||
top_p=top_p, max_tokens=max_tokens, stop=[STOP_STRING])
|
||||
try:
|
||||
resp = client.chat.completions.create(
|
||||
**kwargs, extra_body={"include_stop_str_in_output": True})
|
||||
except Exception:
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
choice = resp.choices[0]
|
||||
return normalize_completion(choice.message.content or "", choice.finish_reason)
|
||||
return fn
|
||||
Reference in New Issue
Block a user