408 lines
16 KiB
Python
408 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .llm_backend import DEFAULT_OLLAMA_HOST, DEFAULT_OLLAMA_MODEL, LLMBackendError, OllamaBackend
|
|
from .permissions import ToolPermissionContext
|
|
from .runtime_tools import RuntimeToolRegistry, render_tool_instructions
|
|
|
|
DEFAULT_AGENT_SYSTEM_PROMPT = (
|
|
'You are a local coding assistant running inside a minimal harness. '
|
|
'Use tools when needed. Respond with exactly one JSON object and no surrounding prose. '
|
|
'Valid outputs are {"type":"tool_call","name":"...","arguments":{...}} or '
|
|
'{"type":"final","content":"..."}. '
|
|
'Do not wrap JSON in markdown fences.'
|
|
)
|
|
MUTATING_TOOL_NAMES = frozenset({'edit_file', 'run_shell_command'})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentRunResult:
|
|
content: str
|
|
turns: int
|
|
tool_calls: tuple[str, ...]
|
|
tool_trace: tuple[dict[str, Any], ...] = ()
|
|
stop_reason: str = 'completed'
|
|
trace_events: tuple[str, ...] = ()
|
|
|
|
|
|
class AgentProtocolError(RuntimeError):
|
|
"""Raised when the model does not follow the JSON tool-call protocol."""
|
|
|
|
|
|
def run_agent_task(
|
|
prompt: str,
|
|
model: str = DEFAULT_OLLAMA_MODEL,
|
|
host: str = DEFAULT_OLLAMA_HOST,
|
|
root: Path | None = None,
|
|
system_prompt: str = DEFAULT_AGENT_SYSTEM_PROMPT,
|
|
max_turns: int = 8,
|
|
permission_context: ToolPermissionContext | None = None,
|
|
trace: bool = False,
|
|
) -> AgentRunResult:
|
|
workspace_root = (root or Path.cwd()).resolve()
|
|
backend = OllamaBackend(model=model, host=host)
|
|
tools = RuntimeToolRegistry(root=workspace_root, permission_context=permission_context)
|
|
messages = build_agent_messages(prompt=prompt, workspace_root=workspace_root, system_prompt=system_prompt)
|
|
tool_calls: list[str] = []
|
|
tool_trace: list[dict[str, Any]] = []
|
|
trace_events: list[str] = []
|
|
tool_required = prompt_requires_tool(prompt)
|
|
edit_required = prompt_requires_edit(prompt)
|
|
required_edit_target = infer_required_edit_target(prompt)
|
|
fallback_used = False
|
|
|
|
for turn in range(1, max_turns + 1):
|
|
response = backend.chat(messages, response_format='json', options={'temperature': 0})
|
|
assistant_content = response.content.strip()
|
|
if trace:
|
|
trace_events.append(f'assistant_raw[{turn}]={assistant_content}')
|
|
messages.append({'role': 'assistant', 'content': assistant_content})
|
|
|
|
try:
|
|
payload = parse_agent_response(assistant_content)
|
|
except AgentProtocolError as error:
|
|
if trace:
|
|
trace_events.append(f'protocol_error[{turn}]={error}')
|
|
fallback = None
|
|
if tool_required and not tool_calls and not fallback_used:
|
|
fallback = infer_tool_call_from_prompt(prompt)
|
|
if fallback is not None:
|
|
fallback_used = True
|
|
if trace:
|
|
trace_events.append(f'fallback_tool[{turn}]={fallback["name"]} {json.dumps(fallback["arguments"], ensure_ascii=True)}')
|
|
_append_tool_result(messages, tools, fallback, tool_calls, tool_trace, trace_events if trace else None, turn)
|
|
continue
|
|
messages.append({'role': 'user', 'content': protocol_error_message(str(error))})
|
|
continue
|
|
|
|
if payload['type'] == 'final':
|
|
if tool_required and not tool_calls:
|
|
if trace:
|
|
trace_events.append(f'tool_required_reject[{turn}]={payload["content"]}')
|
|
fallback = None
|
|
if not fallback_used:
|
|
fallback = infer_tool_call_from_prompt(prompt)
|
|
if fallback is not None:
|
|
fallback_used = True
|
|
if trace:
|
|
trace_events.append(f'fallback_tool[{turn}]={fallback["name"]} {json.dumps(fallback["arguments"], ensure_ascii=True)}')
|
|
_append_tool_result(messages, tools, fallback, tool_calls, tool_trace, trace_events if trace else None, turn)
|
|
continue
|
|
messages.append({'role': 'user', 'content': tool_required_message()})
|
|
continue
|
|
if edit_required and not has_successful_edit(tool_trace, required_edit_target):
|
|
if trace:
|
|
trace_events.append(f'edit_required_reject[{turn}]={payload["content"]}')
|
|
fallback = None
|
|
if not fallback_used:
|
|
fallback = infer_tool_call_from_prompt(prompt, prefer_edit=True)
|
|
if fallback is not None:
|
|
fallback_used = True
|
|
if trace:
|
|
trace_events.append(f'fallback_tool[{turn}]={fallback["name"]} {json.dumps(fallback["arguments"], ensure_ascii=True)}')
|
|
_append_tool_result(messages, tools, fallback, tool_calls, tool_trace, trace_events if trace else None, turn)
|
|
continue
|
|
messages.append({'role': 'user', 'content': edit_required_message()})
|
|
continue
|
|
return AgentRunResult(
|
|
content=payload['content'],
|
|
turns=turn,
|
|
tool_calls=tuple(tool_calls),
|
|
tool_trace=tuple(tool_trace),
|
|
stop_reason='completed',
|
|
trace_events=tuple(trace_events),
|
|
)
|
|
|
|
_append_tool_result(messages, tools, payload, tool_calls, tool_trace, trace_events if trace else None, turn)
|
|
|
|
if trace:
|
|
trace_events.append(f'stopped=max_turns:{max_turns}')
|
|
raise LLMBackendError(f'Agent stopped after reaching max_turns={max_turns}.')
|
|
|
|
|
|
def build_agent_messages(
|
|
prompt: str,
|
|
workspace_root: Path,
|
|
system_prompt: str = DEFAULT_AGENT_SYSTEM_PROMPT,
|
|
) -> list[dict[str, str]]:
|
|
return [
|
|
{
|
|
'role': 'system',
|
|
'content': '\n\n'.join(
|
|
[
|
|
system_prompt,
|
|
f'Workspace root: {workspace_root}',
|
|
render_tool_instructions(),
|
|
'After every tool call you will receive a user message starting with TOOL_RESULT followed by JSON.',
|
|
'If the user asks to inspect workspace files, search code, edit files, or run shell commands, use a tool before giving a final answer.',
|
|
'Example tool call: {"type":"tool_call","name":"read_file","arguments":{"path":"src/main.py","start_line":1,"end_line":40}}',
|
|
'Example final: {"type":"final","content":"I found ..."}',
|
|
]
|
|
),
|
|
},
|
|
{'role': 'user', 'content': prompt},
|
|
]
|
|
|
|
|
|
def parse_agent_response(content: str) -> dict[str, Any]:
|
|
normalized = content.strip()
|
|
if normalized.startswith('```'):
|
|
lines = normalized.splitlines()
|
|
if lines and lines[0].startswith('```'):
|
|
lines = lines[1:]
|
|
if lines and lines[-1].startswith('```'):
|
|
lines = lines[:-1]
|
|
normalized = '\n'.join(lines).strip()
|
|
|
|
try:
|
|
payload = json.loads(normalized)
|
|
except json.JSONDecodeError as error:
|
|
raise AgentProtocolError('Response was not valid JSON.') from error
|
|
|
|
if not isinstance(payload, dict):
|
|
raise AgentProtocolError('Response JSON must be an object.')
|
|
|
|
response_type = payload.get('type')
|
|
if response_type == 'final':
|
|
content_value = payload.get('content')
|
|
if not isinstance(content_value, str) or not content_value:
|
|
raise AgentProtocolError('Final responses must include a non-empty string content field.')
|
|
return {'type': 'final', 'content': content_value}
|
|
|
|
if response_type == 'tool_call':
|
|
name = payload.get('name')
|
|
arguments = payload.get('arguments')
|
|
if not isinstance(name, str) or not name:
|
|
raise AgentProtocolError('Tool calls must include a non-empty string name field.')
|
|
if not isinstance(arguments, dict):
|
|
raise AgentProtocolError('Tool calls must include an arguments object.')
|
|
return {'type': 'tool_call', 'name': name, 'arguments': arguments}
|
|
|
|
raise AgentProtocolError('Response type must be either "tool_call" or "final".')
|
|
|
|
|
|
def tool_result_message(payload: dict[str, Any]) -> str:
|
|
return f'TOOL_RESULT {json.dumps(payload, ensure_ascii=True)}'
|
|
|
|
|
|
def protocol_error_message(error: str) -> str:
|
|
return (
|
|
'Your previous response did not follow the protocol. '
|
|
f'{error} '
|
|
'Respond again with exactly one JSON object.'
|
|
)
|
|
|
|
|
|
def tool_required_message() -> str:
|
|
return (
|
|
'You must use an appropriate tool before answering this request because it depends on workspace contents '
|
|
'or shell output. Respond with exactly one JSON tool_call object.'
|
|
)
|
|
|
|
|
|
def edit_required_message() -> str:
|
|
return (
|
|
'You must make a real workspace change before answering this request. '
|
|
'A read-only inspection is insufficient. Respond with exactly one JSON tool_call object that performs the edit.'
|
|
)
|
|
|
|
|
|
def prompt_requires_tool(prompt: str) -> bool:
|
|
lowered = prompt.lower()
|
|
action_words = ('read', 'search', 'find', 'inspect', 'open', 'edit', 'modify', 'create', 'write', 'run', 'execute', 'quote')
|
|
target_hints = ('src/', '.py', 'file', 'files', 'workspace', 'tree', 'directory', 'lines', 'command', 'shell', 'pwd', 'ls ', 'rg ', 'grep ')
|
|
return any(word in lowered for word in action_words) and any(hint in lowered for hint in target_hints)
|
|
|
|
|
|
def prompt_requires_edit(prompt: str) -> bool:
|
|
lowered = prompt.lower()
|
|
edit_words = ('edit', 'modify', 'change', 'update', 'write', 'create', 'add')
|
|
target_hints = ('.py', '.md', '.txt', 'file', 'workspace', 'train.py', 'src/')
|
|
return any(word in lowered for word in edit_words) and any(hint in lowered for hint in target_hints)
|
|
|
|
|
|
def infer_required_edit_target(prompt: str) -> str | None:
|
|
match = re.search(
|
|
r'(?:modify|edit|change|update|write|create|add(?:\s+to)?)\s+(?P<path>[\w./-]+)',
|
|
prompt,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
if not match:
|
|
return None
|
|
return match.group('path')
|
|
|
|
|
|
def infer_tool_call_from_prompt(prompt: str, prefer_edit: bool = False) -> dict[str, Any] | None:
|
|
if prefer_edit:
|
|
edit_fallback = infer_edit_tool_call_from_prompt(prompt)
|
|
if edit_fallback is not None:
|
|
return edit_fallback
|
|
|
|
edit_fallback = infer_edit_tool_call_from_prompt(prompt)
|
|
if edit_fallback is not None:
|
|
return edit_fallback
|
|
|
|
read_match = re.search(
|
|
r'read\s+(?P<path>[\w./-]+)(?:\s+lines?\s+(?P<start>\d+)\s*(?:to|-)\s*(?P<end>\d+))?',
|
|
prompt,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
if read_match:
|
|
arguments: dict[str, Any] = {'path': read_match.group('path')}
|
|
if read_match.group('start') and read_match.group('end'):
|
|
arguments['start_line'] = int(read_match.group('start'))
|
|
arguments['end_line'] = int(read_match.group('end'))
|
|
return {'name': 'read_file', 'arguments': arguments}
|
|
|
|
search_match = re.search(
|
|
r"search(?:\s+the)?\s+(?P<scope>.+?)\s+for\s+['\"](?P<pattern>[^'\"]+)['\"]",
|
|
prompt,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
if search_match:
|
|
scope = search_match.group('scope').strip().lower()
|
|
path = '.'
|
|
if 'src' in scope:
|
|
path = 'src'
|
|
return {
|
|
'name': 'search_files',
|
|
'arguments': {'path': path, 'pattern': search_match.group('pattern')},
|
|
}
|
|
|
|
quoted_commands = re.findall(r"['\"]([^'\"]+)['\"]", prompt)
|
|
if any(word in prompt.lower() for word in ('run', 'execute', 'shell')) and quoted_commands:
|
|
return {
|
|
'name': 'run_shell_command',
|
|
'arguments': {'command': '\n'.join(quoted_commands)},
|
|
}
|
|
|
|
create_match = re.search(
|
|
r'create\s+(?P<path>[\w./-]+)\s+containing\s+(?P<content>.+)',
|
|
prompt,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
if create_match:
|
|
return {
|
|
'name': 'edit_file',
|
|
'arguments': {
|
|
'path': create_match.group('path'),
|
|
'new_text': create_match.group('content').strip(),
|
|
},
|
|
}
|
|
|
|
if 'git status' in prompt.lower():
|
|
return {'name': 'git_status', 'arguments': {}}
|
|
|
|
if 'git diff' in prompt.lower():
|
|
return {'name': 'git_diff', 'arguments': {}}
|
|
|
|
if 'run tests' in prompt.lower() or 'test suite' in prompt.lower():
|
|
return {'name': 'run_tests', 'arguments': {}}
|
|
|
|
if 'run build' in prompt.lower() or 'build the project' in prompt.lower():
|
|
return {'name': 'run_build', 'arguments': {}}
|
|
|
|
if 'list directory' in prompt.lower() or 'list dir' in prompt.lower():
|
|
return {'name': 'list_dir', 'arguments': {'path': '.'}}
|
|
|
|
return None
|
|
|
|
|
|
def infer_edit_tool_call_from_prompt(prompt: str) -> dict[str, Any] | None:
|
|
modify_match = re.search(
|
|
r'(?:modify|edit|change|update)\s+(?P<path>[\w./-]+)\s+by\s+(?P<details>.+?)(?:\.|\n|$)',
|
|
prompt,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
if modify_match:
|
|
path = modify_match.group('path')
|
|
details = modify_match.group('details').strip().lower()
|
|
if 'comment' in details:
|
|
marker = '# smoke-test comment: '
|
|
command = '\n'.join(
|
|
[
|
|
"python3 - <<'PY'",
|
|
'from pathlib import Path',
|
|
'import re',
|
|
f"target = Path({json.dumps(path)})",
|
|
"text = target.read_text(encoding='utf-8')",
|
|
f"marker = {json.dumps(marker)}",
|
|
"match = re.match(r'^# smoke-test comment: (\\d+)\\n', text)",
|
|
'if match:',
|
|
' current = int(match.group(1))',
|
|
" updated = f'{marker}{current + 1}\\n' + text[match.end():]",
|
|
'else:',
|
|
" updated = f'{marker}1\\n' + text",
|
|
"target.write_text(updated, encoding='utf-8')",
|
|
'PY',
|
|
]
|
|
)
|
|
return {
|
|
'name': 'run_shell_command',
|
|
'arguments': {'command': command, 'timeout_seconds': 20},
|
|
}
|
|
return None
|
|
|
|
|
|
def has_successful_edit(tool_trace: list[dict[str, Any]], required_path: str | None) -> bool:
|
|
for entry in tool_trace:
|
|
if not entry.get('success'):
|
|
continue
|
|
if entry.get('name') not in MUTATING_TOOL_NAMES:
|
|
continue
|
|
metadata = entry.get('metadata')
|
|
if required_path is None:
|
|
return True
|
|
if isinstance(metadata, dict) and metadata.get('path') == required_path:
|
|
return True
|
|
arguments = entry.get('arguments')
|
|
if isinstance(arguments, dict) and arguments.get('path') == required_path:
|
|
return True
|
|
if (
|
|
entry.get('name') == 'run_shell_command'
|
|
and isinstance(arguments, dict)
|
|
and isinstance(arguments.get('command'), str)
|
|
and required_path in arguments['command']
|
|
):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _append_tool_result(
|
|
messages: list[dict[str, str]],
|
|
tools: RuntimeToolRegistry,
|
|
payload: dict[str, Any],
|
|
tool_calls: list[str],
|
|
tool_trace: list[dict[str, Any]],
|
|
trace_events: list[str] | None,
|
|
turn: int,
|
|
) -> None:
|
|
tool_name = payload['name']
|
|
arguments = payload['arguments']
|
|
tool_calls.append(tool_name)
|
|
if trace_events is not None:
|
|
trace_events.append(f'tool_call[{turn}]={tool_name} {json.dumps(arguments, ensure_ascii=True)}')
|
|
tool_result = tools.execute(tool_name, arguments)
|
|
tool_trace.append(
|
|
{
|
|
'turn': turn,
|
|
'name': tool_name,
|
|
'arguments': arguments,
|
|
'success': tool_result.success,
|
|
'output': tool_result.output,
|
|
'metadata': tool_result.metadata,
|
|
'error': tool_result.error,
|
|
}
|
|
)
|
|
if trace_events is not None:
|
|
trace_events.append(
|
|
f'tool_result[{turn}]={tool_name} success={tool_result.success} '
|
|
f'error={tool_result.error!r}'
|
|
)
|
|
messages.append({'role': 'user', 'content': tool_result_message(tool_result.to_payload())})
|