cortex_ai/engine/dsml.py

"""Builders for well-formed DeepSeek-V4 completion text.

The encoding module defines a strict format: ``parse_message_from_completion_text``
validates every token and raises on the slightest deviation. Anything that
produces completions outside the model -- a mock adapter, a replay file, a
conformance test -- should build them with these helpers rather than by hand.

The special tokens are imported from ``encoding_dsv4`` rather than redefined
here. They contain non-ASCII codepoints that do not survive copy-paste, and a
single wrong byte makes every parse fail.
"""

from __future__ import annotations

import json
from typing import Any


def _tokens() -> dict[str, str]:
    try:
        import encoding_dsv4 as enc
    except ImportError as exc:  # pragma: no cover - environment dependent
        raise RuntimeError(
            "The encoding module is required. Add the repository's encoding/ folder "
            "to PYTHONPATH, e.g. PYTHONPATH=/workspace/project/encoding"
        ) from exc
    return {
        "dsml": enc.dsml_token,
        "eos": enc.eos_token,
        "thinking_start": enc.thinking_start_token,
        "thinking_end": enc.thinking_end_token,
        "tool_calls_block": enc.tool_calls_block_name,
    }


def _parameter(name: str, value: Any) -> str:
    dsml = _tokens()["dsml"]
    is_str = isinstance(value, str)
    rendered = value if is_str else json.dumps(value, ensure_ascii=False)
    flag = "true" if is_str else "false"
    return f'<{dsml}parameter name="{name}" string="{flag}">{rendered}</{dsml}parameter>'


def build_tool_call(name: str, arguments: dict[str, Any], reasoning: str = "") -> str:
    """Render one assistant turn that calls a single tool.

    Reproduces the exact byte sequence the parser expects: a leading blank line,
    one parameter per line, the closing invoke tag, and the EOS token.

    In thinking mode the completion opens with the reasoning content and closes
    it with the end tag. The opening `` thinking`` token belongs to the *prompt*,
    so it must not appear here; the parser rejects it inside content.
    """
    t = _tokens()
    lines = [f'<{t["dsml"]}invoke name="{name}">']
    for key, value in arguments.items():
        lines.append(_parameter(key, value))
    body = "\n".join(lines) + f"\n</{t['dsml']}invoke>"
    call = (
        f"\n\n<{t['dsml']}{t['tool_calls_block']}>\n{body}\n"
        f"</{t['dsml']}{t['tool_calls_block']}>"
        f"{t['eos']}"
    )
    if reasoning:
        return f"{reasoning}{t['thinking_end']}{call}"
    return call


def build_answer(content: str, reasoning: str = "") -> str:
    """Render one assistant turn that answers without calling a tool."""
    t = _tokens()
    if reasoning:
        return f"{reasoning}{t['thinking_end']}{content}{t['eos']}"
    return f"{content}{t['eos']}"


def build_tool_result(content: str) -> str:
    """Render a tool result body, matching the encoder's tool_output_template."""
    return f"<tool_result>{content}</tool_result>"