dscli-agent 0.1.0__py3-none-any.whl

dscli/protocol.py

@@ -0,0 +1,633 @@
+from __future__ import annotations
+
+import json
+import re
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Any, Iterable
+
+
+CANONICAL_TOOLS = {
+    "read_file",
+    "write_file",
+    "apply_patch",
+    "list_files",
+    "delete_file",
+    "run_command",
+}
+
+TOOL_ALIASES = {
+    "read": "read_file",
+    "cat": "read_file",
+    "open_file": "read_file",
+    "view_file": "read_file",
+    "write": "write_file",
+    "create_file": "write_file",
+    "create": "write_file",
+    "save_file": "write_file",
+    "overwrite_file": "write_file",
+    "edit": "apply_patch",
+    "edit_file": "apply_patch",
+    "patch": "apply_patch",
+    "replace": "apply_patch",
+    "replace_in_file": "apply_patch",
+    "list": "list_files",
+    "ls": "list_files",
+    "list_dir": "list_files",
+    "list_directory": "list_files",
+    "rm": "delete_file",
+    "remove_file": "delete_file",
+    "bash": "run_command",
+    "shell": "run_command",
+    "exec": "run_command",
+    "execute": "run_command",
+    "run": "run_command",
+    "terminal": "run_command",
+    "finish": "finish",
+    "done": "finish",
+}
+
+
+@dataclass
+class ToolCall:
+    name: str
+    arguments: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Step:
+    thought: str = ""
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    final: str | None = None
+
+    @property
+    def is_final(self) -> bool:
+        return self.final is not None and not self.tool_calls
+
+
+class ProtocolError(ValueError):
+    def __init__(self, message: str, raw: str = "") -> None:
+        super().__init__(message)
+        self.raw = raw
+
+
+def canonical_tool_name(name: Any) -> str:
+    if not isinstance(name, str):
+        return ""
+    cleaned = name.strip().lower().replace("-", "_").replace(" ", "_")
+    cleaned = TOOL_ALIASES.get(cleaned, cleaned)
+    return cleaned
+
+
+def _coerce_args(value: Any) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if isinstance(value, dict):
+        return value
+    if isinstance(value, str):
+        text = value.strip()
+        if not text:
+            return {}
+        try:
+            parsed = json.loads(text)
+        except json.JSONDecodeError:
+            return {"input": value}
+        if isinstance(parsed, dict):
+            return parsed
+        return {"input": parsed}
+    if isinstance(value, list):
+        return {"items": value}
+    return {"input": value}
+
+
+def _normalize_call(obj: Any) -> ToolCall | None:
+    if isinstance(obj, str):
+        name = canonical_tool_name(obj)
+        if name:
+            return ToolCall(name=name, arguments={})
+        return None
+    if not isinstance(obj, dict):
+        return None
+
+    fn = obj.get("function")
+    if isinstance(fn, dict):
+        name = fn.get("name") or fn.get("tool") or fn.get("action")
+        args = fn.get("arguments", fn.get("parameters", fn.get("args", {})))
+    else:
+        name = obj.get("name") or obj.get("tool") or obj.get("action") or obj.get("type")
+        if isinstance(name, str) and name.strip().lower() in {"tool_call", "tool", "function"}:
+            name = obj.get("tool_name") or obj.get("function_name") or obj.get("name")
+        args = obj.get(
+            "arguments",
+            obj.get("args", obj.get("parameters", obj.get("action_input", obj.get("input")))),
+        )
+        if args is None:
+            # Some models flatten arguments into the call object:
+            # {"name":"read_file","path":"x"}.
+            reserved = {
+                "name", "tool", "action", "type", "function", "tool_name",
+                "function_name", "thought", "reasoning",
+            }
+            flattened = {k: v for k, v in obj.items() if k not in reserved}
+            args = flattened if flattened else {}
+    canonical = canonical_tool_name(name)
+    if not canonical:
+        return None
+    return ToolCall(name=canonical, arguments=_coerce_args(args))
+
+
+def _collect_calls(data: dict[str, Any]) -> list[ToolCall]:
+    raw = None
+    for key in ("tool_calls", "toolCalls", "tools", "calls", "actions"):
+        if key in data and data[key] is not None:
+            raw = data[key]
+            break
+    if raw is None:
+        return []
+    if isinstance(raw, dict):
+        raw = [raw]
+    if isinstance(raw, str):
+        try:
+            parsed = json.loads(raw)
+        except json.JSONDecodeError:
+            parsed = [raw]
+        raw = parsed if isinstance(parsed, list) else [parsed]
+    if not isinstance(raw, list):
+        return []
+    calls: list[ToolCall] = []
+    for item in raw:
+        call = _normalize_call(item)
+        if call:
+            calls.append(call)
+    return calls
+
+
+def _coerce_final(value: Any) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (dict, list)):
+        return json.dumps(value, ensure_ascii=False, indent=2)
+    return str(value)
+
+
+def _normalize(data: Any) -> Step | None:
+    if isinstance(data, list):
+        calls = [c for c in (_normalize_call(x) for x in data) if c]
+        calls, finish_text = _finish_to_final(calls)
+        if calls or finish_text is not None:
+            return Step(tool_calls=calls, final=finish_text)
+        return None
+    if not isinstance(data, dict):
+        return None
+
+    thought = data.get("thought") or data.get("reasoning") or data.get("thinking") or ""
+    if not isinstance(thought, str):
+        thought = json.dumps(thought, ensure_ascii=False)
+
+    calls = _collect_calls(data)
+
+    final: Any = None
+    for key in ("final", "final_answer", "answer", "response"):
+        if key in data:
+            final = data[key]
+            break
+
+    dtype = str(data.get("type", "")).strip().lower()
+    if final is None and dtype in {"final", "answer", "result"}:
+        final = data.get("content", data.get("message", data.get("text")))
+    if not calls and dtype in {"tool", "tool_call", "action"}:
+        single = _normalize_call(data)
+        if single:
+            calls = [single]
+    if not calls and isinstance(data.get("action"), str):
+        action_name = canonical_tool_name(data.get("action"))
+        if action_name in CANONICAL_TOOLS or action_name == "finish":
+            calls = [
+                ToolCall(
+                    name=action_name,
+                    arguments=_coerce_args(
+                        data.get("action_input", data.get("arguments", data.get("input", {})))
+                    ),
+                )
+            ]
+    if not calls and isinstance(data.get("tool"), str):
+        single = _normalize_call(data)
+        if single:
+            calls = [single]
+    if not calls and isinstance(data.get("name"), str):
+        single = _normalize_call(data)
+        if single:
+            calls = [single]
+
+    calls, finish_text = _finish_to_final(calls)
+    final_text = _coerce_final(final)
+    if finish_text is not None:
+        final_text = finish_text if final_text is None else final_text
+
+    if not calls and final_text is None:
+        # Some models use a bare {"done": true} at the end.
+        if data.get("done") is True:
+            final_text = _coerce_final(data.get("summary", "")) or "Done."
+        else:
+            return None
+    if final_text is not None:
+        final_text = final_text.strip() if final_text.strip() else final_text
+
+    return Step(thought=thought.strip(), tool_calls=calls, final=final_text)
+
+
+def _finish_to_final(calls: list[ToolCall]) -> tuple[list[ToolCall], str | None]:
+    remaining: list[ToolCall] = []
+    finish_text: str | None = None
+    for call in calls:
+        if call.name == "finish":
+            finish_text = (
+                call.arguments.get("summary")
+                or call.arguments.get("message")
+                or call.arguments.get("final")
+                or call.arguments.get("text")
+                or ""
+            )
+            finish_text = _coerce_final(finish_text) or ""
+        else:
+            remaining.append(call)
+    return remaining, finish_text
+
+
+def _balanced_slices(text: str, max_candidates: int = 40) -> Iterable[str]:
+    """Yield balanced JSON-looking spans starting at every { or [.
+
+    A small model response is short; O(n^2) worst-case is acceptable here.
+    Strings and escapes are respected so braces inside strings are ignored.
+    """
+    n = len(text)
+    produced = 0
+    for start in range(n):
+        if text[start] not in "{[":
+            continue
+        stack: list[str] = []
+        in_string = False
+        escaped = False
+        for i in range(start, n):
+            ch = text[i]
+            if in_string:
+                if escaped:
+                    escaped = False
+                elif ch == "\\":
+                    escaped = True
+                elif ch == '"':
+                    in_string = False
+                continue
+            if ch == '"':
+                in_string = True
+            elif ch in "{[":
+                stack.append(ch)
+            elif ch in "}]":
+                if not stack:
+                    break
+                opener = stack.pop()
+                if (opener == "{" and ch != "}") or (opener == "[" and ch != "]"):
+                    break
+                if not stack:
+                    yield text[start : i + 1]
+                    produced += 1
+                    break
+        if produced >= max_candidates:
+            break
+
+
+_FENCE_RE = re.compile(r"```[ \t]*(?:json|JSON)?[ \t]*\r?\n?(.*?)```", re.DOTALL)
+
+
+def _json_candidates(raw: str) -> list[str]:
+    candidates: list[str] = []
+    seen: set[str] = set()
+
+    def add(text: str) -> None:
+        text = text.strip()
+        if text and text not in seen:
+            seen.add(text)
+            candidates.append(text)
+
+    add(raw)
+    for match in _FENCE_RE.finditer(raw):
+        add(match.group(1))
+    for span in _balanced_slices(raw):
+        add(span)
+    # A fenced block may itself contain prose around a JSON object; the
+    # balanced scan above already covers that.
+    return candidates
+
+
+def _repair_invalid_escapes(text: str) -> str:
+    repaired: list[str] = []
+    in_string = False
+    index = 0
+    while index < len(text):
+        char = text[index]
+        if char == '"':
+            in_string = not in_string
+            repaired.append(char)
+            index += 1
+            continue
+        if char != "\\" or not in_string or index + 1 >= len(text):
+            repaired.append(char)
+            index += 1
+            continue
+
+        escape = text[index + 1]
+        if escape in {'"', "\\", "/", "b", "f", "n", "r", "t"}:
+            repaired.append(text[index : index + 2])
+            index += 2
+        elif (
+            escape == "u"
+            and index + 6 <= len(text)
+            and all(c in "0123456789abcdefABCDEF" for c in text[index + 2 : index + 6])
+        ):
+            repaired.append(text[index : index + 6])
+            index += 6
+        elif escape == "u":
+            repaired.append(char)
+            index += 1
+        else:
+            repaired.extend(("\\", "\\", escape))
+            index += 2
+    return "".join(repaired)
+
+
+def _repair_string_content(text: str, max_attempts: int = 300) -> str | None:
+    """Escape quotes and backslashes that break otherwise valid JSON strings.
+
+    Models often paste code into ``content`` or ``command`` values without
+    escaping its double quotes. Each parse error tells us which quote the
+    parser just treated as the end of a string, so escape that one and retry.
+    """
+    pending = deque([text])
+    seen = {text}
+    while pending and len(seen) <= max_attempts:
+        current = pending.popleft()
+        try:
+            json.loads(current)
+        except json.JSONDecodeError as exc:
+            repairs: list[str] = []
+            if exc.msg == "Expecting ',' delimiter":
+                positions: list[int] = []
+                anchor = exc.pos - 1
+                while anchor >= 0 and current[anchor].isspace():
+                    anchor -= 1
+                if anchor >= 0 and current[anchor] == '"':
+                    positions.append(anchor)
+                if exc.pos < len(current) and current[exc.pos] == '"':
+                    positions.append(exc.pos)
+                for pos in positions:
+                    repairs.append(current[:pos] + "\\" + current[pos:])
+            elif exc.msg == "Invalid \\escape":
+                if exc.pos < len(current) and current[exc.pos] == "\\":
+                    repairs.append(current[: exc.pos] + "\\" + current[exc.pos:])
+            for repaired in repairs:
+                if repaired not in seen:
+                    seen.add(repaired)
+                    pending.append(repaired)
+        else:
+            return current
+    return None
+
+
+def _try_loads(text: str) -> Any:
+    bases = [text]
+    stripped = text.strip()
+    if stripped.lower().startswith("json\n"):
+        bases.append(stripped[5:])
+    attempts: list[str] = []
+    for base in bases:
+        repaired = _repair_invalid_escapes(base)
+        attempts.extend(
+            [
+                base,
+                re.sub(r",\s*([}\]])", r"\1", base),
+                repaired,
+                re.sub(r",\s*([}\]])", r"\1", repaired),
+            ]
+        )
+    last_error: Exception | None = None
+    for attempt in attempts:
+        try:
+            return json.loads(attempt)
+        except (json.JSONDecodeError, TypeError) as exc:
+            last_error = exc
+    for attempt in dict.fromkeys(attempts):
+        repaired = _repair_string_content(attempt)
+        if repaired is None:
+            continue
+        for candidate in (repaired, re.sub(r",\s*([}\]])", r"\1", repaired)):
+            try:
+                return json.loads(candidate)
+            except (json.JSONDecodeError, TypeError) as exc:
+                last_error = exc
+    raise last_error or ValueError("empty JSON")
+
+
+def _complete_json_string(text: str, start: int) -> tuple[str, int] | None:
+    escaped = False
+    for index in range(start + 1, len(text)):
+        char = text[index]
+        if escaped:
+            escaped = False
+        elif char == "\\":
+            escaped = True
+        elif char == '"':
+            try:
+                value = json.loads(text[start : index + 1])
+            except json.JSONDecodeError:
+                return None
+            return value, index + 1
+    return None
+
+
+def _json_value_end(text: str, start: int) -> int | None:
+    depth = 0
+    in_string = False
+    escaped = False
+    for index in range(start, len(text)):
+        char = text[index]
+        if in_string:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == '"':
+                in_string = False
+            continue
+        if char == '"':
+            in_string = True
+        elif char in "[{":
+            depth += 1
+        elif char in "]}":
+            if depth == 0:
+                return index if char == "}" else None
+            depth -= 1
+        elif char == "," and depth == 0:
+            return index
+    return None
+
+
+def _decode_final_prefix(text: str, start: int) -> str:
+    decoded: list[str] = []
+    escapes = {
+        '"': '"',
+        "\\": "\\",
+        "/": "/",
+        "b": "\b",
+        "f": "\f",
+        "n": "\n",
+        "r": "\r",
+        "t": "\t",
+    }
+    index = start + 1
+    while index < len(text):
+        char = text[index]
+        if char == '"':
+            break
+        if char != "\\":
+            decoded.append(char)
+            index += 1
+            continue
+        if index + 1 >= len(text):
+            break
+        escape = text[index + 1]
+        if escape in escapes:
+            decoded.append(escapes[escape])
+            index += 2
+            continue
+        if escape != "u":
+            decoded.extend(("\\", escape))
+            index += 2
+            continue
+        if index + 6 > len(text):
+            break
+        digits = text[index + 2 : index + 6]
+        if any(c not in "0123456789abcdefABCDEF" for c in digits):
+            break
+        codepoint = int(digits, 16)
+        next_index = index + 6
+        if 0xD800 <= codepoint <= 0xDFFF:
+            if next_index + 6 > len(text) or text[next_index : next_index + 2] != "\\u":
+                break
+            low_digits = text[next_index + 2 : next_index + 6]
+            if any(c not in "0123456789abcdefABCDEF" for c in low_digits):
+                break
+            low = int(low_digits, 16)
+            if not (0xD800 <= codepoint <= 0xDBFF and 0xDC00 <= low <= 0xDFFF):
+                break
+            codepoint = 0x10000 + ((codepoint - 0xD800) << 10) + (low - 0xDC00)
+            next_index += 6
+        decoded.append(chr(codepoint))
+        index = next_index
+    return "".join(decoded)
+
+
+_FENCE_OPEN_RE = re.compile(r"```[ \t]*(?:json)?[ \t]*\r?\n", re.IGNORECASE)
+
+
+def extract_final_prefix(raw: str) -> str | None:
+    text = raw.lstrip()
+    fence = _FENCE_OPEN_RE.match(text)
+    if fence:
+        text = text[fence.end() :].lstrip()
+    if not text.startswith("{"):
+        return None
+    index = 1
+    tool_calls_empty = False
+    while index < len(text):
+        while index < len(text) and text[index].isspace():
+            index += 1
+        if index >= len(text) or text[index] != '"':
+            return None
+        key_result = _complete_json_string(text, index)
+        if key_result is None:
+            return None
+        key, index = key_result
+        while index < len(text) and text[index].isspace():
+            index += 1
+        if index >= len(text) or text[index] != ":":
+            return None
+        index += 1
+        while index < len(text) and text[index].isspace():
+            index += 1
+        if index >= len(text):
+            return None
+        if key == "final":
+            if not tool_calls_empty or text[index] != '"':
+                return None
+            return _decode_final_prefix(text, index)
+
+        value_end = _json_value_end(text, index)
+        if value_end is None:
+            return None
+        if key == "tool_calls":
+            tool_calls_empty = re.fullmatch(r"\[\s*\]", text[index:value_end]) is not None
+        index = value_end
+        if index >= len(text) or text[index] != ",":
+            return None
+        index += 1
+    return None
+
+
+def _score(data: Any) -> int:
+    if isinstance(data, dict):
+        score = 0
+        keys = set(data)
+        if "tool_calls" in keys or "toolCalls" in keys:
+            score += 120
+        if "final" in keys or "final_answer" in keys:
+            score += 110
+        if "answer" in keys or "response" in keys:
+            score += 60
+        if "name" in keys and ("arguments" in keys or "args" in keys):
+            score += 80
+        if "tool" in keys:
+            score += 50
+        if "action" in keys:
+            score += 20
+        if "thought" in keys or "reasoning" in keys:
+            score += 10
+        return score
+    if isinstance(data, list) and data and all(isinstance(x, dict) for x in data):
+        return 40
+    return 0
+
+
+def parse_step(raw: str) -> Step:
+    if not raw or not raw.strip():
+        raise ProtocolError("The model returned an empty response.", raw)
+
+    best: tuple[int, int, Step] | None = None
+    errors: list[str] = []
+    for candidate in _json_candidates(raw):
+        try:
+            data = _try_loads(candidate)
+        except Exception as exc:
+            errors.append(f"{type(exc).__name__}: {exc}")
+            continue
+        step = _normalize(data)
+        if step is None:
+            continue
+        score = _score(data)
+        rank = (score, len(candidate), step)
+        if best is None or rank[:2] > best[:2]:
+            best = rank
+    if best is None:
+        detail = errors[-1] if errors else "no complete JSON object was found"
+        raise ProtocolError(
+            "The model's response did not contain a valid dscli JSON envelope "
+            f"({detail}).",
+            raw,
+        )
+    step = best[2]
+    if not step.tool_calls and step.final is None:
+        raise ProtocolError(
+            "The JSON envelope has neither tool_calls nor a final answer.", raw
+        )
+    return step