Yif-Yang
673 lines
20 KiB
Python
673 lines
20 KiB
Python
"""Codex CLI backend for ReflACT."""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import mimetypes
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
from urllib.parse import unquote, urlparse
|
|
|
|
from skillopt.model.backend_config import get_codex_exec_config
|
|
from skillopt.model.common import (
|
|
CompatAssistantMessage,
|
|
CompatToolCall,
|
|
CompatToolFunction,
|
|
TokenTracker,
|
|
)
|
|
|
|
|
|
CODEX_BIN = os.environ.get("CODEX_CLI_BIN", "codex")
|
|
CODEX_PROFILE = os.environ.get("CODEX_PROFILE", "review")
|
|
CODEX_SANDBOX_MODE = os.environ.get("CODEX_SANDBOX_MODE", "read-only")
|
|
|
|
OPTIMIZER_DEPLOYMENT = os.environ.get("OPTIMIZER_DEPLOYMENT", "gpt-4o")
|
|
TARGET_DEPLOYMENT = os.environ.get("TARGET_DEPLOYMENT", "gpt-4o")
|
|
|
|
REASONING_EFFORT: str | None = None
|
|
tracker = TokenTracker()
|
|
|
|
|
|
def _default_working_directory() -> str:
|
|
return os.environ.get("CODEX_WORKING_DIRECTORY", os.getcwd())
|
|
|
|
|
|
def _parse_data_uri(url: str) -> tuple[bytes, str]:
|
|
header, data = url.split(",", 1)
|
|
mime = header[5:].split(";", 1)[0] or "image/png"
|
|
return base64.b64decode(data), mime
|
|
|
|
|
|
def _content_to_text(
|
|
content: Any,
|
|
attachments: list[dict[str, Any]],
|
|
*,
|
|
image_counter: int,
|
|
) -> tuple[str, int]:
|
|
if isinstance(content, str):
|
|
return content, image_counter
|
|
|
|
if not isinstance(content, list):
|
|
return str(content), image_counter
|
|
|
|
parts: list[str] = []
|
|
for item in content:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
item_type = item.get("type")
|
|
if item_type == "text":
|
|
parts.append(str(item.get("text", "")))
|
|
continue
|
|
if item_type != "image_url":
|
|
continue
|
|
|
|
image_counter += 1
|
|
label = f"[Attached image {image_counter}]"
|
|
parts.append(label)
|
|
|
|
image_url = item.get("image_url", {}) or {}
|
|
url = str(image_url.get("url", "") or "")
|
|
if not url:
|
|
continue
|
|
if url.startswith("data:") and ";base64," in url:
|
|
data, mime = _parse_data_uri(url)
|
|
attachments.append({"bytes": data, "mime": mime})
|
|
continue
|
|
if url.startswith("file://"):
|
|
parsed = urlparse(url)
|
|
path = unquote(parsed.path)
|
|
if path:
|
|
attachments.append({"path": path})
|
|
continue
|
|
if os.path.exists(url):
|
|
attachments.append({"path": url})
|
|
|
|
return "".join(parts), image_counter
|
|
|
|
|
|
def _simplify_tool_schemas(tools: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
|
simplified: list[dict[str, Any]] = []
|
|
for tool in tools or []:
|
|
function = tool.get("function", tool)
|
|
simplified.append(
|
|
{
|
|
"name": function.get("name", ""),
|
|
"description": function.get("description", ""),
|
|
"parameters": function.get("parameters", {}),
|
|
}
|
|
)
|
|
return simplified
|
|
|
|
|
|
def _build_prompt_from_messages(
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
structured_output: bool = False,
|
|
) -> tuple[str, list[dict[str, Any]]]:
|
|
system_parts: list[str] = []
|
|
history_parts: list[str] = []
|
|
attachments: list[dict[str, Any]] = []
|
|
image_counter = 0
|
|
|
|
def _history_line(label: str, body: str) -> str:
|
|
stripped = body.strip()
|
|
if not stripped:
|
|
return f"- {label}:"
|
|
indented = stripped.replace("\n", "\n ")
|
|
return f"- {label}: {indented}"
|
|
|
|
for message in messages:
|
|
role = str(message.get("role", "user"))
|
|
text, image_counter = _content_to_text(
|
|
message.get("content", ""),
|
|
attachments,
|
|
image_counter=image_counter,
|
|
)
|
|
|
|
if role == "system":
|
|
if text.strip():
|
|
system_parts.append(text.strip())
|
|
continue
|
|
|
|
if role == "assistant":
|
|
block = _history_line("Assistant", text)
|
|
tool_calls = message.get("tool_calls") or []
|
|
if tool_calls:
|
|
simplified_calls = []
|
|
for tool_call in tool_calls:
|
|
function = tool_call.get("function", {}) or {}
|
|
simplified_calls.append(
|
|
{
|
|
"name": function.get("name", ""),
|
|
"arguments": function.get("arguments", "{}"),
|
|
}
|
|
)
|
|
block += (
|
|
"\n Compatibility tool requests:\n"
|
|
+ json.dumps(simplified_calls, ensure_ascii=False, indent=2)
|
|
)
|
|
history_parts.append(block)
|
|
continue
|
|
|
|
if role == "tool":
|
|
tool_call_id = str(message.get("tool_call_id", "") or "")
|
|
label = f"Tool result (tool_call_id={tool_call_id})"
|
|
history_parts.append(_history_line(label, text))
|
|
continue
|
|
|
|
history_parts.append(_history_line(role.capitalize(), text))
|
|
|
|
prompt_parts: list[str] = []
|
|
|
|
system_text = "\n\n".join(part for part in system_parts if part).strip()
|
|
if system_text:
|
|
prompt_parts.append(system_text)
|
|
|
|
if tools:
|
|
simplified_tools = _simplify_tool_schemas(tools)
|
|
prompt_parts.append(
|
|
"Available compatibility tools:\n"
|
|
+ json.dumps(simplified_tools, ensure_ascii=False, indent=2)
|
|
)
|
|
prompt_parts.append(
|
|
"Do not execute these tools yourself. If you need one, request it in "
|
|
"`tool_calls`. Each `arguments` field must be a JSON string."
|
|
)
|
|
|
|
if tool_choice == "required":
|
|
prompt_parts.append(
|
|
"Tool choice policy: you must request at least one compatibility tool."
|
|
)
|
|
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
|
function = tool_choice.get("function", {}) or {}
|
|
prompt_parts.append(
|
|
"Tool choice policy: you must request the compatibility tool "
|
|
f"`{function.get('name', '')}`."
|
|
)
|
|
|
|
history_text = "\n".join(part for part in history_parts if part).strip()
|
|
if history_text:
|
|
prompt_parts.append("History:\n" + history_text)
|
|
|
|
if structured_output:
|
|
prompt_parts.append("Return only JSON matching the provided schema.")
|
|
if tools:
|
|
prompt_parts.append(
|
|
"Set `content` to the assistant-visible reply. Set `tool_calls` to "
|
|
"an empty array when no tool is needed."
|
|
)
|
|
else:
|
|
prompt_parts.append("Answer the latest user request.")
|
|
|
|
return "\n\n".join(prompt_parts), attachments
|
|
|
|
|
|
def _assistant_message_schema() -> dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"content": {"type": "string"},
|
|
"tool_calls": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"arguments": {"type": "string"},
|
|
},
|
|
"required": ["name", "arguments"],
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
},
|
|
"required": ["content", "tool_calls"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
|
|
def _materialize_attachments(
|
|
attachments: list[dict[str, Any]],
|
|
temp_dir: str,
|
|
) -> list[str]:
|
|
image_paths: list[str] = []
|
|
for index, attachment in enumerate(attachments, 1):
|
|
path = attachment.get("path")
|
|
if path:
|
|
image_paths.append(str(path))
|
|
continue
|
|
|
|
mime = str(attachment.get("mime", "image/png"))
|
|
suffix = mimetypes.guess_extension(mime) or ".png"
|
|
image_path = os.path.join(temp_dir, f"image_{index}{suffix}")
|
|
with open(image_path, "wb") as f:
|
|
f.write(attachment.get("bytes", b""))
|
|
image_paths.append(image_path)
|
|
return image_paths
|
|
|
|
|
|
def _usage_from_event(usage: dict[str, Any] | None) -> dict[str, int]:
|
|
usage = usage or {}
|
|
prompt_tokens = int(usage.get("input_tokens", 0) or 0)
|
|
completion_tokens = int(usage.get("output_tokens", 0) or 0)
|
|
return {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
}
|
|
|
|
|
|
def _extract_error(stdout: str, stderr: str) -> str:
|
|
for raw_line in reversed(stdout.splitlines()):
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
payload = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if payload.get("type") == "turn.failed":
|
|
error = payload.get("error", {}) or {}
|
|
return str(error.get("message", "") or "Codex turn failed")
|
|
if payload.get("type") == "error":
|
|
return str(payload.get("message", "") or "Codex execution failed")
|
|
return stderr.strip() or stdout.strip() or "Codex execution failed"
|
|
|
|
|
|
def _run_codex_exec(
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
attachments: list[dict[str, Any]],
|
|
output_schema: dict[str, Any] | None,
|
|
timeout: int | None,
|
|
) -> tuple[str, dict[str, int]]:
|
|
with tempfile.TemporaryDirectory(prefix="skillopt_codex_") as temp_dir:
|
|
config = get_codex_exec_config()
|
|
output_path = os.path.join(temp_dir, "last_message.txt")
|
|
image_paths = _materialize_attachments(attachments, temp_dir)
|
|
profile = str(config.get("profile") or os.environ.get("CODEX_PROFILE", "")).strip()
|
|
reasoning_effort = str(REASONING_EFFORT or config.get("reasoning_effort") or "").strip()
|
|
|
|
command = [
|
|
str(config.get("path") or CODEX_BIN),
|
|
"exec",
|
|
"--json",
|
|
"--ephemeral",
|
|
"-c",
|
|
f"approval_policy={json.dumps(str(config.get('approval_policy') or 'never'))}",
|
|
"--sandbox",
|
|
str(config.get("sandbox") or CODEX_SANDBOX_MODE),
|
|
"--skip-git-repo-check",
|
|
"--cd",
|
|
_default_working_directory(),
|
|
"--model",
|
|
model,
|
|
"--output-last-message",
|
|
output_path,
|
|
]
|
|
|
|
if profile:
|
|
command.extend(["--profile", profile])
|
|
|
|
if reasoning_effort and reasoning_effort != "none":
|
|
command.extend(["-c", f"model_reasoning_effort={json.dumps(reasoning_effort)}"])
|
|
|
|
schema_path = None
|
|
if output_schema is not None:
|
|
schema_path = os.path.join(temp_dir, "schema.json")
|
|
with open(schema_path, "w", encoding="utf-8") as f:
|
|
json.dump(output_schema, f, ensure_ascii=False)
|
|
command.extend(["--output-schema", schema_path])
|
|
|
|
for image_path in image_paths:
|
|
command.extend(["--image", image_path])
|
|
|
|
command.append("-")
|
|
|
|
proc = subprocess.run(
|
|
command,
|
|
input=prompt,
|
|
text=True,
|
|
encoding="utf-8",
|
|
errors="replace",
|
|
capture_output=True,
|
|
timeout=timeout,
|
|
check=False,
|
|
)
|
|
|
|
usage_info = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
|
fallback_text = ""
|
|
for raw_line in proc.stdout.splitlines():
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
payload = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if payload.get("type") == "item.completed":
|
|
item = payload.get("item", {}) or {}
|
|
if item.get("type") == "agent_message":
|
|
fallback_text = str(item.get("text", "") or fallback_text)
|
|
if payload.get("type") == "turn.completed":
|
|
usage_info = _usage_from_event(payload.get("usage"))
|
|
|
|
last_message = ""
|
|
if os.path.exists(output_path):
|
|
with open(output_path, encoding="utf-8") as f:
|
|
last_message = f.read().strip()
|
|
if not last_message:
|
|
last_message = fallback_text.strip()
|
|
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(_extract_error(proc.stdout, proc.stderr))
|
|
if not last_message:
|
|
raise RuntimeError("Codex returned an empty final message")
|
|
return last_message, usage_info
|
|
|
|
|
|
def _tool_name_from_choice(tool_choice: str | dict[str, Any] | None) -> str | None:
|
|
if not isinstance(tool_choice, dict):
|
|
return None
|
|
if tool_choice.get("type") != "function":
|
|
return None
|
|
function = tool_choice.get("function", {}) or {}
|
|
return str(function.get("name", "") or "") or None
|
|
|
|
|
|
def _compat_message_from_payload(
|
|
payload: dict[str, Any],
|
|
*,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
) -> CompatAssistantMessage:
|
|
content = str(payload.get("content", "") or "")
|
|
tool_calls: list[CompatToolCall] = []
|
|
for index, raw_tool_call in enumerate(payload.get("tool_calls", []) or [], 1):
|
|
if not isinstance(raw_tool_call, dict):
|
|
continue
|
|
name = str(raw_tool_call.get("name", "") or "")
|
|
arguments = raw_tool_call.get("arguments", "{}")
|
|
if not isinstance(arguments, str):
|
|
arguments = json.dumps(arguments, ensure_ascii=False)
|
|
tool_calls.append(
|
|
CompatToolCall(
|
|
id=f"tool_{index}_{uuid.uuid4().hex[:12]}",
|
|
function=CompatToolFunction(name=name, arguments=arguments),
|
|
)
|
|
)
|
|
|
|
if tool_choice == "required" and not tool_calls:
|
|
raise RuntimeError("Codex response did not request a tool under tool_choice='required'")
|
|
|
|
required_name = _tool_name_from_choice(tool_choice)
|
|
if required_name and all(
|
|
tool_call.function.name != required_name for tool_call in tool_calls
|
|
):
|
|
raise RuntimeError(
|
|
f"Codex response did not request the required tool {required_name!r}"
|
|
)
|
|
|
|
return CompatAssistantMessage(content=content, tool_calls=tool_calls)
|
|
|
|
|
|
def _chat_messages_impl(
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int,
|
|
retries: int,
|
|
stage: str,
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
return_message: bool = False,
|
|
timeout: int | None = None,
|
|
) -> tuple[Any, dict[str, int]]:
|
|
del max_completion_tokens
|
|
last_err = None
|
|
structured_output = bool(tools) or return_message
|
|
|
|
for attempt in range(retries):
|
|
try:
|
|
prompt, attachments = _build_prompt_from_messages(
|
|
messages,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
structured_output=structured_output,
|
|
)
|
|
raw_text, usage_info = _run_codex_exec(
|
|
model=model,
|
|
prompt=prompt,
|
|
attachments=attachments,
|
|
output_schema=_assistant_message_schema() if structured_output else None,
|
|
timeout=timeout,
|
|
)
|
|
tracker.record(
|
|
stage,
|
|
usage_info["prompt_tokens"],
|
|
usage_info["completion_tokens"],
|
|
)
|
|
|
|
if not structured_output:
|
|
return raw_text, usage_info
|
|
|
|
payload = json.loads(raw_text)
|
|
compat = _compat_message_from_payload(payload, tool_choice=tool_choice)
|
|
return (compat if return_message else compat.content), usage_info
|
|
except subprocess.TimeoutExpired as exc:
|
|
last_err = RuntimeError(f"Codex CLI timed out after {timeout}s") if timeout else exc
|
|
except Exception as exc: # noqa: BLE001
|
|
last_err = exc
|
|
time.sleep(min(2 ** attempt, 30))
|
|
|
|
raise RuntimeError(f"Codex call failed after {retries} retries: {last_err}")
|
|
|
|
|
|
def chat_with_model(
|
|
model: str,
|
|
system: str,
|
|
user: str,
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "custom",
|
|
timeout: int | None = None,
|
|
) -> tuple[str, dict[str, int]]:
|
|
messages = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
]
|
|
return _chat_messages_impl(
|
|
model,
|
|
messages,
|
|
max_completion_tokens,
|
|
retries,
|
|
stage,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_messages_with_model(
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "custom",
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
return_message: bool = False,
|
|
timeout: int | None = None,
|
|
) -> tuple[Any, dict[str, int]]:
|
|
return _chat_messages_impl(
|
|
model,
|
|
messages,
|
|
max_completion_tokens,
|
|
retries,
|
|
stage,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
return_message=return_message,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_optimizer(
|
|
system: str,
|
|
user: str,
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "optimizer",
|
|
timeout: int | None = None,
|
|
) -> tuple[str, dict[str, int]]:
|
|
return chat_with_model(
|
|
model=OPTIMIZER_DEPLOYMENT,
|
|
system=system,
|
|
user=user,
|
|
max_completion_tokens=max_completion_tokens,
|
|
retries=retries,
|
|
stage=stage,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_with_deployment(
|
|
deployment: str,
|
|
system: str,
|
|
user: str,
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "custom",
|
|
timeout: int | None = None,
|
|
) -> tuple[str, dict[str, int]]:
|
|
return chat_with_model(
|
|
model=deployment,
|
|
system=system,
|
|
user=user,
|
|
max_completion_tokens=max_completion_tokens,
|
|
retries=retries,
|
|
stage=stage,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_target(
|
|
system: str,
|
|
user: str,
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "target",
|
|
timeout: int | None = None,
|
|
) -> tuple[str, dict[str, int]]:
|
|
return chat_with_model(
|
|
model=TARGET_DEPLOYMENT,
|
|
system=system,
|
|
user=user,
|
|
max_completion_tokens=max_completion_tokens,
|
|
retries=retries,
|
|
stage=stage,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_optimizer_messages(
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "optimizer",
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
return_message: bool = False,
|
|
timeout: int | None = None,
|
|
) -> tuple[Any, dict[str, int]]:
|
|
return _chat_messages_impl(
|
|
OPTIMIZER_DEPLOYMENT,
|
|
messages,
|
|
max_completion_tokens,
|
|
retries,
|
|
stage,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
return_message=return_message,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_messages_with_deployment(
|
|
deployment: str,
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "custom",
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
return_message: bool = False,
|
|
timeout: int | None = None,
|
|
) -> tuple[Any, dict[str, int]]:
|
|
return _chat_messages_impl(
|
|
deployment,
|
|
messages,
|
|
max_completion_tokens,
|
|
retries,
|
|
stage,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
return_message=return_message,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def chat_target_messages(
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int = 16384,
|
|
retries: int = 5,
|
|
stage: str = "target",
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
return_message: bool = False,
|
|
timeout: int | None = None,
|
|
) -> tuple[Any, dict[str, int]]:
|
|
return _chat_messages_impl(
|
|
TARGET_DEPLOYMENT,
|
|
messages,
|
|
max_completion_tokens,
|
|
retries,
|
|
stage,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
return_message=return_message,
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def get_token_summary() -> dict[str, dict[str, int]]:
|
|
return tracker.summary()
|
|
|
|
|
|
def reset_token_tracker() -> None:
|
|
tracker.reset()
|
|
|
|
|
|
def set_target_deployment(deployment: str) -> None:
|
|
global TARGET_DEPLOYMENT
|
|
TARGET_DEPLOYMENT = deployment
|
|
os.environ["TARGET_DEPLOYMENT"] = deployment
|
|
|
|
|
|
def set_reasoning_effort(effort: str | None) -> None:
|
|
global REASONING_EFFORT
|
|
REASONING_EFFORT = effort if effort else None
|
|
|
|
|
|
def set_optimizer_deployment(deployment: str) -> None:
|
|
global OPTIMIZER_DEPLOYMENT
|
|
OPTIMIZER_DEPLOYMENT = deployment
|
|
os.environ["OPTIMIZER_DEPLOYMENT"] = deployment
|