payload/app/scripts/dg_openai_tool_loop.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path
from typing import Any


DEFAULT_BASE_URL = "http://127.0.0.1:4100/v1"
DEFAULT_MODEL = "diffusiongemma-local"
DEFAULT_TOOL_MANIFEST_URL = "http://127.0.0.1:8090/v1/agent/tool_manifest"
DEFAULT_TOOL_RUNTIME_URL = "http://127.0.0.1:8090/v1/agent/tool"
REPO_TOOL_NAMES = {
    "dg_repo_status",
    "dg_git_diff",
    "dg_list_files",
    "dg_read_file",
    "dg_search",
}
OSS_REPO_TOOL_NAMES = {
    "dg_repo_pack",
    "dg_repo_map",
    "dg_ast_grep",
    "dg_code_outline",
}
READ_ONLY_TOOL_NAMES = REPO_TOOL_NAMES | OSS_REPO_TOOL_NAMES | {
    "dg_context",
    "dg_rag_context",
    "dg_session_artifact",
    "dg_agent_run_artifact",
}
EXPLICIT_READ_ONLY_TOOL_NAMES = {"dg_agent_run_artifact"}
EXPLICIT_TOOL_NAMES = {"dg_agent_run_artifact"}

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")


def request_json(method: str, url: str, payload: dict[str, Any] | None = None, timeout: int = 60) -> dict[str, Any]:
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        method=method,
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer dummy",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        text = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{method} {url} failed with HTTP {exc.code}: {text}") from exc


def tool_schema_name(tool: dict[str, Any]) -> str:
    if isinstance(tool.get("function"), dict):
        return str(tool["function"].get("name") or "")
    return str(tool.get("name") or "")


def load_tools(
    manifest_url: str,
    timeout: int,
    *,
    include_execute_command: bool,
    include_tools: list[str],
    exclude_tools: list[str],
    read_only: bool,
) -> list[dict[str, Any]]:
    manifest = request_json("GET", manifest_url, timeout=timeout)
    tools = manifest.get("openai_chat_completions", {}).get("tools", [])
    if not isinstance(tools, list) or not tools:
        raise RuntimeError(f"no chat tools in manifest: {manifest_url}")
    if not include_execute_command:
        tools = [tool for tool in tools if tool_schema_name(tool) != "execute_command"]
    if not include_tools:
        tools = [tool for tool in tools if tool_schema_name(tool) not in EXPLICIT_TOOL_NAMES]
    if read_only:
        allowed_tools = set(READ_ONLY_TOOL_NAMES)
        if not include_tools:
            allowed_tools -= EXPLICIT_READ_ONLY_TOOL_NAMES
        tools = [tool for tool in tools if tool_schema_name(tool) in allowed_tools]
    if include_tools:
        include_set = set(include_tools)
        tools = [tool for tool in tools if tool_schema_name(tool) in include_set]
    if exclude_tools:
        exclude_set = set(exclude_tools)
        tools = [tool for tool in tools if tool_schema_name(tool) not in exclude_set]
    if not tools:
        raise RuntimeError(f"no usable chat tools in manifest: {manifest_url}")
    return tools


def tool_call_to_dict(tool_call: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": tool_call.get("id", ""),
        "type": tool_call.get("type", "function"),
        "function": {
            "name": tool_call.get("function", {}).get("name", ""),
            "arguments": tool_call.get("function", {}).get("arguments", "{}"),
        },
    }


def patch_tool_call_repo(tool_call: dict[str, Any], repo: str) -> dict[str, Any]:
    if not repo:
        return tool_call
    function = tool_call.get("function")
    if not isinstance(function, dict):
        return tool_call
    name = str(function.get("name") or "")
    if name not in REPO_TOOL_NAMES and name not in OSS_REPO_TOOL_NAMES and name not in {"dg_session", "dg_context", "dg_rag_context"}:
        return tool_call
    raw_args = function.get("arguments") or "{}"
    try:
        parsed = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
    except Exception:
        return tool_call
    if not isinstance(parsed, dict):
        return tool_call
    if not parsed.get("repo") or parsed.get("repo") == ".":
        parsed["repo"] = repo
        function["arguments"] = json.dumps(parsed, ensure_ascii=False)
    return tool_call


def task_file_hints(task: str, limit: int = 8) -> list[str]:
    hints: list[str] = []
    for match in re.findall(r"[\w./\\-]+\.[A-Za-z0-9_]+", task):
        text = match.replace("\\", "/").strip("./")
        if text and text not in hints:
            hints.append(text)
    return hints[:limit]


def task_search_hint(task: str) -> str:
    quoted = re.search(r"[`'\"]([^`'\"]{2,120})[`'\"]", task)
    if quoted:
        return quoted.group(1).strip()
    match = re.search(r"\b(?:search|grep|find|where)\s+(?:for\s+)?(.{2,120})", task, flags=re.I)
    if match:
        return match.group(1).strip(" .")
    return task[:120].strip()


def select_tool(tools: list[dict[str, Any]], names: set[str]) -> dict[str, Any] | None:
    for tool in tools:
        if tool_schema_name(tool) in names:
            return tool
    return None


def deterministic_tool_args(name: str, task: str, repo: str) -> dict[str, Any] | None:
    task_lower = task.lower()
    files = task_file_hints(task)
    if name in REPO_TOOL_NAMES or name in OSS_REPO_TOOL_NAMES or name in {"dg_context", "dg_rag_context"}:
        if not repo:
            return None
    if name == "dg_repo_status":
        return {"repo": repo}
    if name == "dg_git_diff":
        args: dict[str, Any] = {"repo": repo, "stat": "stat" in task_lower}
        if "cached" in task_lower or "staged" in task_lower:
            args["cached"] = True
        if files:
            args["files"] = files
        return args
    if name == "dg_list_files":
        args = {"repo": repo, "limit": 200}
        glob_match = re.search(r"(\*\.[A-Za-z0-9_]+)", task)
        if glob_match:
            args["globs"] = [glob_match.group(1)]
        return args
    if name == "dg_read_file":
        if not files:
            return None
        return {"repo": repo, "path": files[0], "start_line": 1, "max_lines": 160}
    if name == "dg_search":
        args = {"repo": repo, "query": task_search_hint(task), "max_matches": 80}
        glob_match = re.search(r"(\*\.[A-Za-z0-9_]+)", task)
        if glob_match:
            args["globs"] = [glob_match.group(1)]
        return args
    if name == "dg_repo_pack":
        return {"repo": repo, "style": "markdown", "max_chars": 20000}
    if name == "dg_repo_map":
        args = {"repo": repo, "map_tokens": 2048, "map_only": True, "max_chars": 20000}
        if files:
            args["paths"] = files
        return args
    if name == "dg_ast_grep":
        args = {"repo": repo, "pattern": task_search_hint(task), "max_matches": 80}
        if "python" in task_lower or "*.py" in task_lower or any(file.endswith(".py") for file in files):
            args["lang"] = "python"
        if files:
            args["paths"] = files
        return args
    if name == "dg_code_outline":
        args = {"repo": repo, "items": "auto", "view": "auto", "max_items": 200}
        if "python" in task_lower or "*.py" in task_lower or any(file.endswith(".py") for file in files):
            args["lang"] = "python"
        if files:
            args["paths"] = files
        return args
    if name == "dg_context":
        args = {"repo": repo, "task": task, "format": "json", "max_files": 3, "max_snippet_chars": 1200}
        if files:
            args["files"] = files
        return args
    if name == "dg_rag_context":
        return {"repo": repo, "task": task, "max_context_chars": 1200, "max_files": 4, "max_tokens": 128}
    return None


def deterministic_tool_call(args: argparse.Namespace, tools: list[dict[str, Any]]) -> dict[str, Any] | None:
    available = {tool_schema_name(tool) for tool in tools}
    task_lower = args.task.lower()
    explicit = [name for name in args.tool if name in available]
    for name in explicit:
        tool_args = deterministic_tool_args(name, args.task, args.repo)
        if tool_args is not None:
            return build_tool_call(name, tool_args)

    if not args.read_only:
        return None
    candidates: list[str] = []
    if any(word in task_lower for word in ("git status", "repo status", "working tree", "untracked")):
        candidates.append("dg_repo_status")
    if any(word in task_lower for word in ("git diff", "diff stat", "working diff", "staged diff", "cached diff")):
        candidates.append("dg_git_diff")
    if any(word in task_lower for word in ("list files", "show files", "file list", "which files")):
        candidates.append("dg_list_files")
    if any(word in task_lower for word in ("read file", "show file", "open file", "cat ")) and task_file_hints(args.task):
        candidates.append("dg_read_file")
    if any(word in task_lower for word in ("code outline", "outline", "symbols", "symbol list", "structure")):
        candidates.append("dg_code_outline")
    if any(word in task_lower for word in ("repo map", "repository map", "aider map")):
        candidates.append("dg_repo_map")
    if any(word in task_lower for word in ("ast-grep", "ast grep", "structural search")):
        candidates.append("dg_ast_grep")
    if any(word in task_lower for word in ("search", "grep", "find", "inspect", "explain", "where", "context", "summarize")):
        candidates.extend(["dg_rag_context", "dg_context", "dg_search"])
    candidates.extend(["dg_rag_context", "dg_context"])

    for name in candidates:
        if name not in available:
            continue
        tool_args = deterministic_tool_args(name, args.task, args.repo)
        if tool_args is not None:
            return build_tool_call(name, tool_args)
    return None


def build_tool_call(name: str, args: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": "call_" + uuid.uuid4().hex[:24],
        "type": "function",
        "function": {
            "name": name,
            "arguments": json.dumps(args, ensure_ascii=False),
        },
    }


def runtime_content(runtime: dict[str, Any]) -> str:
    result = runtime.get("result", {}) if isinstance(runtime.get("result"), dict) else {}
    content = str(result.get("content") or result.get("text") or result.get("stdout") or "")
    if not content and isinstance(result.get("files"), list):
        content = "\n".join(str(item) for item in result["files"])
    if not content and isinstance(result.get("matches"), list):
        content = "\n".join(str(item) for item in result["matches"])
    if not content:
        content = json.dumps(result or runtime, ensure_ascii=False, indent=2)
    return content


def deterministic_final(runtime: dict[str, Any], tool_call: dict[str, Any]) -> str:
    name = tool_call.get("function", {}).get("name", "tool")
    result = runtime.get("result", {}) if isinstance(runtime.get("result"), dict) else {}
    status = result.get("status", "success" if runtime.get("ok") else "failed")
    return f"Deterministic {name} result ({status}):\n\n{runtime_content(runtime)}"


def run_loop(args: argparse.Namespace) -> dict[str, Any]:
    tools = load_tools(
        args.tool_manifest_url,
        args.timeout,
        include_execute_command=args.include_execute_command,
        include_tools=args.tool,
        exclude_tools=args.exclude_tool,
        read_only=args.read_only,
    )
    base_url = args.base_url.rstrip("/")
    chat_url = f"{base_url}/chat/completions"
    messages: list[dict[str, Any]] = []
    if args.system:
        messages.append({"role": "system", "content": args.system})
    messages.append({"role": "user", "content": args.task})

    events: list[dict[str, Any]] = []
    final_content = ""
    started = time.time()

    if args.deterministic_first:
        direct_call = deterministic_tool_call(args, tools)
        if direct_call is not None:
            direct_call = patch_tool_call_repo(direct_call, args.repo)
            assistant_message = {"role": "assistant", "content": None, "tool_calls": [direct_call]}
            messages.append(assistant_message)
            runtime = request_json("POST", args.tool_runtime_url, {"tool_call": direct_call}, timeout=args.timeout)
            events.append({"step": 0, "kind": "tool_runtime", "deterministic": True, "tool_call": direct_call, "runtime": runtime})
            tool_response = runtime.get("tool_response")
            if not isinstance(tool_response, dict):
                tool_response = {
                    "role": "tool",
                    "tool_call_id": direct_call.get("id", ""),
                    "content": json.dumps(runtime, ensure_ascii=False),
                }
            messages.append(tool_response)
            final_content = deterministic_final(runtime, direct_call)
            elapsed = time.time() - started
            return {
                "status": "success" if runtime.get("ok") else "failed",
                "route": "deterministic_tool_runtime",
                "elapsed_seconds": elapsed,
                "base_url": base_url,
                "model": args.model,
                "tool_manifest_url": args.tool_manifest_url,
                "tool_runtime_url": args.tool_runtime_url,
                "tool_names": [tool_schema_name(tool) for tool in tools],
                "steps": 0,
                "final_content": final_content,
                "messages": messages,
                "events": events,
            }

    for step in range(1, args.max_steps + 1):
        payload = {
            "model": args.model,
            "messages": messages,
            "tools": tools,
            "tool_choice": "auto",
            "temperature": args.temperature,
            "max_tokens": args.max_tokens,
        }
        response = request_json("POST", chat_url, payload, timeout=args.timeout)
        choice = (response.get("choices") or [{}])[0]
        message = choice.get("message") or {}
        tool_calls = message.get("tool_calls") or []
        events.append({"step": step, "kind": "chat", "finish_reason": choice.get("finish_reason"), "message": message})

        if not tool_calls:
            final_content = str(message.get("content") or "")
            if final_content:
                messages.append({"role": "assistant", "content": final_content})
            break

        patched_calls = [patch_tool_call_repo(tool_call_to_dict(item), args.repo) for item in tool_calls if isinstance(item, dict)]
        assistant_message = {
            "role": "assistant",
            "content": message.get("content"),
            "tool_calls": patched_calls,
        }
        messages.append(assistant_message)

        for item in assistant_message["tool_calls"]:
            runtime_payload = {"tool_call": item}
            runtime = request_json("POST", args.tool_runtime_url, runtime_payload, timeout=args.timeout)
            events.append({"step": step, "kind": "tool_runtime", "tool_call": item, "runtime": runtime})
            tool_response = runtime.get("tool_response")
            if not isinstance(tool_response, dict):
                tool_response = {
                    "role": "tool",
                    "tool_call_id": item.get("id", ""),
                    "content": json.dumps(runtime, ensure_ascii=False),
                }
            messages.append(tool_response)

        if args.stop_after_tool:
            break

    elapsed = time.time() - started
    return {
        "status": "success",
        "route": "openai_chat_tool_loop",
        "elapsed_seconds": elapsed,
        "base_url": base_url,
        "model": args.model,
        "tool_manifest_url": args.tool_manifest_url,
        "tool_runtime_url": args.tool_runtime_url,
        "tool_names": [tool_schema_name(tool) for tool in tools],
        "steps": len([event for event in events if event["kind"] == "chat"]),
        "final_content": final_content,
        "messages": messages,
        "events": events,
    }


def print_text(report: dict[str, Any]) -> None:
    print(f"Status: {report['status']}")
    print(f"Model: {report['model']}")
    print(f"Steps: {report['steps']}")
    print(f"Tools: {', '.join(report.get('tool_names') or [])}")
    print(f"Elapsed: {report['elapsed_seconds']:.3f}s")
    print("")
    for event in report["events"]:
        if event["kind"] == "chat":
            print(f"[chat step {event['step']}] finish={event.get('finish_reason')}")
            tool_calls = event.get("message", {}).get("tool_calls") or []
            for call in tool_calls:
                fn = call.get("function", {})
                print(f"  tool_call {call.get('id', '')}: {fn.get('name', '')} {fn.get('arguments', '')[:300]}")
            content = event.get("message", {}).get("content")
            if content:
                print(str(content).strip())
        elif event["kind"] == "tool_runtime":
            runtime = event.get("runtime", {})
            result = runtime.get("result", {}) if isinstance(runtime.get("result"), dict) else {}
            print(f"[tool step {event['step']}] {runtime.get('tool')} ok={runtime.get('ok')} status={result.get('status', '')}")
            content = str(result.get("content") or result.get("text") or "")
            if content:
                print(content[:1200].rstrip())
    if report.get("final_content"):
        print("\nFinal:")
        print(report["final_content"].strip())


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run a small OpenAI-compatible DG tool loop.")
    parser.add_argument("--task", required=True)
    parser.add_argument("--repo", default="")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--tool-manifest-url", default=DEFAULT_TOOL_MANIFEST_URL)
    parser.add_argument("--tool-runtime-url", default=DEFAULT_TOOL_RUNTIME_URL)
    parser.add_argument("--system", default="Use DG-specific tools when helpful.")
    parser.add_argument("--max-steps", type=int, default=2)
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--timeout", type=int, default=120)
    parser.add_argument("--stop-after-tool", action="store_true")
    parser.add_argument("--no-deterministic-first", dest="deterministic_first", action="store_false", help="Disable direct rule-based read-only tool routing before model calls.")
    parser.add_argument("--include-execute-command", action="store_true", help="Include the legacy execute_command schema. Off by default to prefer DG-specific tools.")
    parser.add_argument("--tool", action="append", default=[], help="Limit the manifest to this tool name. Repeatable.")
    parser.add_argument("--exclude-tool", action="append", default=[], help="Remove this tool name from the manifest. Repeatable.")
    parser.add_argument("--read-only", action="store_true", help="Expose only read-only repo/context/artifact tools.")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--out", default="")
    parser.set_defaults(deterministic_first=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    report = run_loop(args)
    if args.out:
        path = Path(args.out).resolve()
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if args.json:
        print(json.dumps(report, ensure_ascii=False, indent=2))
    else:
        print_text(report)
        if args.out:
            print(f"\nTranscript: {Path(args.out).resolve()}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())