#!/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())