kimi-cli 0.44__py3-none-any.whl → 0.78__py3-none-any.whl

kimi_cli/flow/__init__.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Literal
+
+from kosong.message import ContentPart
+
+FlowNodeKind = Literal["begin", "end", "task", "decision"]
+
+
+class PromptFlowError(ValueError):
+    """Base error for prompt flow parsing/validation."""
+
+
+class PromptFlowParseError(PromptFlowError):
+    """Raised when prompt flow parsing fails."""
+
+
+class PromptFlowValidationError(PromptFlowError):
+    """Raised when a flowchart fails validation."""
+
+
+@dataclass(frozen=True, slots=True)
+class FlowNode:
+    id: str
+    label: str | list[ContentPart]
+    kind: FlowNodeKind
+
+
+@dataclass(frozen=True, slots=True)
+class FlowEdge:
+    src: str
+    dst: str
+    label: str | None
+
+
+@dataclass(slots=True)
+class PromptFlow:
+    nodes: dict[str, FlowNode]
+    outgoing: dict[str, list[FlowEdge]]
+    begin_id: str
+    end_id: str
+
+
+_CHOICE_RE = re.compile(r"<choice>([^<]*)</choice>")
+
+
+def parse_choice(text: str) -> str | None:
+    matches = _CHOICE_RE.findall(text or "")
+    if not matches:
+        return None
+    return matches[-1].strip()
+
+
+def validate_flow(
+    nodes: dict[str, FlowNode],
+    outgoing: dict[str, list[FlowEdge]],
+) -> tuple[str, str]:
+    begin_ids = [node.id for node in nodes.values() if node.kind == "begin"]
+    end_ids = [node.id for node in nodes.values() if node.kind == "end"]
+
+    if len(begin_ids) != 1:
+        raise PromptFlowValidationError(f"Expected exactly one BEGIN node, found {len(begin_ids)}")
+    if len(end_ids) != 1:
+        raise PromptFlowValidationError(f"Expected exactly one END node, found {len(end_ids)}")
+
+    begin_id = begin_ids[0]
+    end_id = end_ids[0]
+
+    reachable: set[str] = set()
+    queue: list[str] = [begin_id]
+    while queue:
+        node_id = queue.pop()
+        if node_id in reachable:
+            continue
+        reachable.add(node_id)
+        for edge in outgoing.get(node_id, []):
+            if edge.dst not in reachable:
+                queue.append(edge.dst)
+
+    for node in nodes.values():
+        if node.id not in reachable:
+            continue
+        edges = outgoing.get(node.id, [])
+        if node.kind == "begin":
+            if len(edges) != 1:
+                raise PromptFlowValidationError("BEGIN node must have exactly one outgoing edge")
+            continue
+        if node.kind == "end":
+            if edges:
+                raise PromptFlowValidationError("END node must not have outgoing edges")
+            continue
+        if node.kind == "decision":
+            if not edges:
+                raise PromptFlowValidationError(
+                    f'Decision node "{node.id}" must have outgoing edges'
+                )
+            labels: list[str] = []
+            for edge in edges:
+                if edge.label is None or not edge.label.strip():
+                    raise PromptFlowValidationError(
+                        f'Decision node "{node.id}" has an unlabeled edge'
+                    )
+                labels.append(edge.label)
+            if len(set(labels)) != len(labels):
+                raise PromptFlowValidationError(
+                    f'Decision node "{node.id}" has duplicate edge labels'
+                )
+            continue
+        if len(edges) != 1:
+            raise PromptFlowValidationError(f'Node "{node.id}" must have exactly one outgoing edge')
+
+    if end_id not in reachable:
+        raise PromptFlowValidationError("END node is not reachable from BEGIN")
+
+    return begin_id, end_id

kimi_cli/flow/d2.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from . import (
+    FlowEdge,
+    FlowNode,
+    FlowNodeKind,
+    PromptFlow,
+    PromptFlowParseError,
+    validate_flow,
+)
+
+_NODE_ID_RE = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_./-]*")
+_PROPERTY_SEGMENTS = {
+    "shape",
+    "style",
+    "label",
+    "link",
+    "icon",
+    "near",
+    "width",
+    "height",
+    "direction",
+    "grid-rows",
+    "grid-columns",
+    "grid-gap",
+    "font-size",
+    "font-family",
+    "font-color",
+    "stroke",
+    "fill",
+    "opacity",
+    "padding",
+    "border-radius",
+    "shadow",
+    "sketch",
+    "animated",
+    "multiple",
+    "constraint",
+    "tooltip",
+}
+
+
+@dataclass(frozen=True, slots=True)
+class _NodeDef:
+    node: FlowNode
+    explicit: bool
+
+
+def parse_d2_flowchart(text: str) -> PromptFlow:
+    nodes: dict[str, _NodeDef] = {}
+    outgoing: dict[str, list[FlowEdge]] = {}
+
+    for line_no, statement in _iter_top_level_statements(text):
+        if _has_unquoted_token(statement, "->"):
+            _parse_edge_statement(statement, line_no, nodes, outgoing)
+        else:
+            _parse_node_statement(statement, line_no, nodes)
+
+    flow_nodes = {node_id: node_def.node for node_id, node_def in nodes.items()}
+    for node_id in flow_nodes:
+        outgoing.setdefault(node_id, [])
+
+    flow_nodes = _infer_decision_nodes(flow_nodes, outgoing)
+    begin_id, end_id = validate_flow(flow_nodes, outgoing)
+    return PromptFlow(nodes=flow_nodes, outgoing=outgoing, begin_id=begin_id, end_id=end_id)
+
+
+def _iter_top_level_statements(text: str) -> Iterable[tuple[int, str]]:
+    text = text.replace("\r\n", "\n").replace("\r", "\n")
+    brace_depth = 0
+    in_single = False
+    in_double = False
+    escape = False
+    drop_line = False
+    buf: list[str] = []
+    line_no = 1
+    stmt_line = 1
+    i = 0
+
+    while i < len(text):
+        ch = text[i]
+        next_ch = text[i + 1] if i + 1 < len(text) else ""
+
+        if ch == "\\" and next_ch == "\n":
+            i += 2
+            line_no += 1
+            continue
+
+        if ch == "\n":
+            if brace_depth == 0 and not in_single and not in_double and not drop_line:
+                statement = "".join(buf).strip()
+                if statement:
+                    yield stmt_line, statement
+            buf = []
+            drop_line = False
+            stmt_line = line_no + 1
+            line_no += 1
+            i += 1
+            continue
+
+        if not in_single and not in_double:
+            if ch == "#":
+                while i < len(text) and text[i] != "\n":
+                    i += 1
+                continue
+            if ch == "{":
+                if brace_depth == 0:
+                    statement = "".join(buf).strip()
+                    if statement:
+                        yield stmt_line, statement
+                    drop_line = True
+                    buf.clear()
+                brace_depth += 1
+                i += 1
+                continue
+            if ch == "}" and brace_depth > 0:
+                brace_depth -= 1
+                i += 1
+                continue
+            if ch == "}" and brace_depth == 0:
+                raise PromptFlowParseError(_line_error(line_no, "Unmatched '}'"))
+
+        if ch == "'" and not in_double and not escape:
+            in_single = not in_single
+        elif ch == '"' and not in_single and not escape:
+            in_double = not in_double
+
+        if escape:
+            escape = False
+        elif ch == "\\" and (in_single or in_double):
+            escape = True
+
+        if brace_depth == 0 and not drop_line:
+            buf.append(ch)
+
+        i += 1
+
+    if brace_depth != 0:
+        raise PromptFlowParseError(_line_error(line_no, "Unclosed '{' block"))
+    if in_single or in_double:
+        raise PromptFlowParseError(_line_error(line_no, "Unclosed string"))
+
+    statement = "".join(buf).strip()
+    if statement:
+        yield stmt_line, statement
+
+
+def _has_unquoted_token(text: str, token: str) -> bool:
+    parts = _split_on_token(text, token)
+    return len(parts) > 1
+
+
+def _parse_edge_statement(
+    statement: str,
+    line_no: int,
+    nodes: dict[str, _NodeDef],
+    outgoing: dict[str, list[FlowEdge]],
+) -> None:
+    parts = _split_on_token(statement, "->")
+    if len(parts) < 2:
+        raise PromptFlowParseError(_line_error(line_no, "Expected edge arrow"))
+
+    last_part = parts[-1]
+    target_text, edge_label = _split_unquoted_once(last_part, ":")
+    parts[-1] = target_text
+
+    node_ids: list[str] = []
+    for idx, part in enumerate(parts):
+        node_id = _parse_node_id(part, line_no, allow_inline_label=(idx < len(parts) - 1))
+        node_ids.append(node_id)
+
+    if any(_is_property_path(node_id) for node_id in node_ids):
+        return
+    if len(node_ids) < 2:
+        raise PromptFlowParseError(_line_error(line_no, "Edge must have at least two nodes"))
+
+    label = _parse_label(edge_label, line_no) if edge_label is not None else None
+    for idx in range(len(node_ids) - 1):
+        edge = FlowEdge(
+            src=node_ids[idx],
+            dst=node_ids[idx + 1],
+            label=label if idx == len(node_ids) - 2 else None,
+        )
+        outgoing.setdefault(edge.src, []).append(edge)
+        outgoing.setdefault(edge.dst, [])
+
+    for node_id in node_ids:
+        _add_node(nodes, node_id=node_id, label=None, explicit=False, line_no=line_no)
+
+
+def _parse_node_statement(statement: str, line_no: int, nodes: dict[str, _NodeDef]) -> None:
+    node_text, label_text = _split_unquoted_once(statement, ":")
+    if label_text is not None and _is_property_path(node_text):
+        return
+    node_id = _parse_node_id(node_text, line_no, allow_inline_label=False)
+    label = None
+    explicit = False
+    if label_text is not None and not label_text.strip():
+        return
+    if label_text is not None:
+        label = _parse_label(label_text, line_no)
+        explicit = True
+    _add_node(nodes, node_id=node_id, label=label, explicit=explicit, line_no=line_no)
+
+
+def _parse_node_id(text: str, line_no: int, *, allow_inline_label: bool) -> str:
+    cleaned = text.strip()
+    if allow_inline_label and ":" in cleaned:
+        cleaned = _split_unquoted_once(cleaned, ":")[0].strip()
+    if not cleaned:
+        raise PromptFlowParseError(_line_error(line_no, "Expected node id"))
+    match = _NODE_ID_RE.fullmatch(cleaned)
+    if not match:
+        raise PromptFlowParseError(_line_error(line_no, f'Invalid node id "{cleaned}"'))
+    return match.group(0)
+
+
+def _is_property_path(node_id: str) -> bool:
+    if "." not in node_id:
+        return False
+    parts = [part for part in node_id.split(".") if part]
+    for part in parts[1:]:
+        if part in _PROPERTY_SEGMENTS or part.startswith("style"):
+            return True
+    return parts[-1] in _PROPERTY_SEGMENTS
+
+
+def _parse_label(text: str, line_no: int) -> str:
+    label = text.strip()
+    if not label:
+        raise PromptFlowParseError(_line_error(line_no, "Label cannot be empty"))
+    if label[0] in {"'", '"'}:
+        return _parse_quoted_label(label, line_no)
+    return label
+
+
+def _parse_quoted_label(text: str, line_no: int) -> str:
+    quote = text[0]
+    buf: list[str] = []
+    escape = False
+    i = 1
+    while i < len(text):
+        ch = text[i]
+        if escape:
+            buf.append(ch)
+            escape = False
+            i += 1
+            continue
+        if ch == "\\":
+            escape = True
+            i += 1
+            continue
+        if ch == quote:
+            trailing = text[i + 1 :].strip()
+            if trailing:
+                raise PromptFlowParseError(_line_error(line_no, "Unexpected trailing content"))
+            return "".join(buf)
+        buf.append(ch)
+        i += 1
+    raise PromptFlowParseError(_line_error(line_no, "Unclosed quoted label"))
+
+
+def _split_on_token(text: str, token: str) -> list[str]:
+    parts: list[str] = []
+    buf: list[str] = []
+    in_single = False
+    in_double = False
+    escape = False
+    i = 0
+
+    while i < len(text):
+        if not in_single and not in_double and text.startswith(token, i):
+            parts.append("".join(buf).strip())
+            buf = []
+            i += len(token)
+            continue
+        ch = text[i]
+        if escape:
+            escape = False
+        elif ch == "\\" and (in_single or in_double):
+            escape = True
+        elif ch == "'" and not in_double:
+            in_single = not in_single
+        elif ch == '"' and not in_single:
+            in_double = not in_double
+        buf.append(ch)
+        i += 1
+
+    if in_single or in_double:
+        raise PromptFlowParseError("Unclosed string in statement")
+    parts.append("".join(buf).strip())
+    return parts
+
+
+def _split_unquoted_once(text: str, token: str) -> tuple[str, str | None]:
+    in_single = False
+    in_double = False
+    escape = False
+    for idx, ch in enumerate(text):
+        if escape:
+            escape = False
+            continue
+        if ch == "\\" and (in_single or in_double):
+            escape = True
+            continue
+        if ch == "'" and not in_double:
+            in_single = not in_single
+            continue
+        if ch == '"' and not in_single:
+            in_double = not in_double
+            continue
+        if ch == token and not in_single and not in_double:
+            return text[:idx].strip(), text[idx + 1 :].strip()
+    return text.strip(), None
+
+
+def _add_node(
+    nodes: dict[str, _NodeDef],
+    *,
+    node_id: str,
+    label: str | None,
+    explicit: bool,
+    line_no: int,
+) -> FlowNode:
+    label = label if label is not None else node_id
+    label_norm = label.strip().lower()
+    if not label:
+        raise PromptFlowParseError(_line_error(line_no, "Node label cannot be empty"))
+
+    kind: FlowNodeKind = "task"
+    if label_norm == "begin":
+        kind = "begin"
+    elif label_norm == "end":
+        kind = "end"
+
+    node = FlowNode(id=node_id, label=label, kind=kind)
+    existing = nodes.get(node_id)
+    if existing is None:
+        nodes[node_id] = _NodeDef(node=node, explicit=explicit)
+        return node
+
+    if existing.node == node:
+        return existing.node
+
+    if not explicit and existing.explicit:
+        return existing.node
+
+    if explicit and not existing.explicit:
+        nodes[node_id] = _NodeDef(node=node, explicit=True)
+        return node
+
+    raise PromptFlowParseError(_line_error(line_no, f'Conflicting definition for node "{node_id}"'))
+
+
+def _infer_decision_nodes(
+    nodes: dict[str, FlowNode],
+    outgoing: dict[str, list[FlowEdge]],
+) -> dict[str, FlowNode]:
+    updated: dict[str, FlowNode] = {}
+    for node_id, node in nodes.items():
+        kind = node.kind
+        if kind == "task" and len(outgoing.get(node_id, [])) > 1:
+            kind = "decision"
+        if kind != node.kind:
+            updated[node_id] = FlowNode(id=node.id, label=node.label, kind=kind)
+        else:
+            updated[node_id] = node
+    return updated
+
+
+def _line_error(line_no: int, message: str) -> str:
+    return f"Line {line_no}: {message}"

kimi_cli/flow/mermaid.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from . import (
+    FlowEdge,
+    FlowNode,
+    FlowNodeKind,
+    PromptFlow,
+    PromptFlowParseError,
+    validate_flow,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _NodeSpec:
+    node_id: str
+    label: str | None
+    shape: str | None
+
+
+@dataclass(slots=True)
+class _NodeDef:
+    node: FlowNode
+    explicit: bool
+
+
+_NODE_ID_RE = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_-]*")
+_HEADER_RE = re.compile(r"^(flowchart|graph)\b", re.IGNORECASE)
+
+_SHAPES = {
+    "[": ("square", "]"),
+    "(": ("paren", ")"),
+    "{": ("curly", "}"),
+}
+
+
+def parse_mermaid_flowchart(text: str) -> PromptFlow:
+    nodes: dict[str, _NodeDef] = {}
+    outgoing: dict[str, list[FlowEdge]] = {}
+
+    for line_no, raw_line in enumerate(text.splitlines(), start=1):
+        line = raw_line.strip()
+        if not line or line.startswith("%%"):
+            continue
+        if _HEADER_RE.match(line):
+            continue
+        if "-->" in line:
+            src_spec, label, dst_spec = _parse_edge_line(line, line_no)
+            src_node = _add_node(nodes, src_spec, line_no)
+            dst_node = _add_node(nodes, dst_spec, line_no)
+            edge = FlowEdge(src=src_node.id, dst=dst_node.id, label=label)
+            outgoing.setdefault(edge.src, []).append(edge)
+            outgoing.setdefault(edge.dst, [])
+            continue
+
+        node_spec, idx = _parse_node_token(line, 0, line_no)
+        idx = _skip_ws(line, idx)
+        if idx != len(line):
+            raise PromptFlowParseError(_line_error(line_no, "Unexpected trailing content"))
+        _add_node(nodes, node_spec, line_no)
+
+    flow_nodes = {node_id: node_def.node for node_id, node_def in nodes.items()}
+    for node_id in flow_nodes:
+        outgoing.setdefault(node_id, [])
+
+    begin_id, end_id = validate_flow(flow_nodes, outgoing)
+    return PromptFlow(nodes=flow_nodes, outgoing=outgoing, begin_id=begin_id, end_id=end_id)
+
+
+def _parse_edge_line(line: str, line_no: int) -> tuple[_NodeSpec, str | None, _NodeSpec]:
+    src_spec, idx = _parse_node_token(line, 0, line_no)
+    idx = _skip_ws(line, idx)
+    if line.startswith("-->", idx):
+        idx += 3
+        idx = _skip_ws(line, idx)
+        label = None
+        if idx < len(line) and line[idx] == "|":
+            label, idx = _parse_pipe_label(line, idx, line_no)
+            idx = _skip_ws(line, idx)
+        dst_spec, idx = _parse_node_token(line, idx, line_no)
+        idx = _skip_ws(line, idx)
+        if idx != len(line):
+            raise PromptFlowParseError(_line_error(line_no, "Unexpected trailing content"))
+        return src_spec, label, dst_spec
+
+    if line.startswith("--", idx):
+        idx += 2
+        arrow_idx = line.find("-->", idx)
+        if arrow_idx == -1:
+            raise PromptFlowParseError(_line_error(line_no, "Expected '-->' to end edge label"))
+        label = line[idx:arrow_idx].strip()
+        if not label:
+            raise PromptFlowParseError(_line_error(line_no, "Edge label cannot be empty"))
+        idx = arrow_idx + 3
+        idx = _skip_ws(line, idx)
+        dst_spec, idx = _parse_node_token(line, idx, line_no)
+        idx = _skip_ws(line, idx)
+        if idx != len(line):
+            raise PromptFlowParseError(_line_error(line_no, "Unexpected trailing content"))
+        return src_spec, label, dst_spec
+
+    raise PromptFlowParseError(_line_error(line_no, "Expected edge arrow"))
+
+
+def _parse_node_token(line: str, idx: int, line_no: int) -> tuple[_NodeSpec, int]:
+    match = _NODE_ID_RE.match(line, idx)
+    if not match:
+        raise PromptFlowParseError(_line_error(line_no, "Expected node id"))
+    node_id = match.group(0)
+    idx = match.end()
+
+    if idx >= len(line) or line[idx] not in _SHAPES:
+        return _NodeSpec(node_id=node_id, label=None, shape=None), idx
+
+    shape, close_char = _SHAPES[line[idx]]
+    idx += 1
+    label, idx = _parse_label(line, idx, close_char, line_no)
+    return _NodeSpec(node_id=node_id, label=label, shape=shape), idx
+
+
+def _parse_label(line: str, idx: int, close_char: str, line_no: int) -> tuple[str, int]:
+    if idx >= len(line):
+        raise PromptFlowParseError(_line_error(line_no, "Expected node label"))
+    if close_char == ")" and line[idx] == "[":
+        label, idx = _parse_label(line, idx + 1, "]", line_no)
+        while idx < len(line) and line[idx].isspace():
+            idx += 1
+        if idx >= len(line) or line[idx] != ")":
+            raise PromptFlowParseError(_line_error(line_no, "Unclosed node label"))
+        return label, idx + 1
+    if line[idx] == '"':
+        idx += 1
+        buf: list[str] = []
+        while idx < len(line):
+            ch = line[idx]
+            if ch == '"':
+                idx += 1
+                while idx < len(line) and line[idx].isspace():
+                    idx += 1
+                if idx >= len(line) or line[idx] != close_char:
+                    raise PromptFlowParseError(_line_error(line_no, "Unclosed node label"))
+                return "".join(buf), idx + 1
+            if ch == "\\" and idx + 1 < len(line):
+                buf.append(line[idx + 1])
+                idx += 2
+                continue
+            buf.append(ch)
+            idx += 1
+        raise PromptFlowParseError(_line_error(line_no, "Unclosed quoted label"))
+
+    end = line.find(close_char, idx)
+    if end == -1:
+        raise PromptFlowParseError(_line_error(line_no, "Unclosed node label"))
+    label = line[idx:end].strip()
+    if not label:
+        raise PromptFlowParseError(_line_error(line_no, "Node label cannot be empty"))
+    return label, end + 1
+
+
+def _parse_pipe_label(line: str, idx: int, line_no: int) -> tuple[str, int]:
+    if line[idx] != "|":
+        raise PromptFlowParseError(_line_error(line_no, "Expected '|' for edge label"))
+    end = line.find("|", idx + 1)
+    if end == -1:
+        raise PromptFlowParseError(_line_error(line_no, "Unclosed edge label"))
+    label = line[idx + 1 : end].strip()
+    if not label:
+        raise PromptFlowParseError(_line_error(line_no, "Edge label cannot be empty"))
+    return label, end + 1
+
+
+def _skip_ws(line: str, idx: int) -> int:
+    while idx < len(line) and line[idx].isspace():
+        idx += 1
+    return idx
+
+
+def _add_node(nodes: dict[str, _NodeDef], spec: _NodeSpec, line_no: int) -> FlowNode:
+    label = spec.label if spec.label is not None else spec.node_id
+    label_norm = label.strip().lower()
+    if not label:
+        raise PromptFlowParseError(_line_error(line_no, "Node label cannot be empty"))
+
+    kind: FlowNodeKind = "task"
+    if spec.shape == "curly":
+        kind = "decision"
+    if label_norm == "begin":
+        kind = "begin"
+    elif label_norm == "end":
+        kind = "end"
+
+    node = FlowNode(id=spec.node_id, label=label, kind=kind)
+    explicit = spec.label is not None
+
+    existing = nodes.get(spec.node_id)
+    if existing is None:
+        nodes[spec.node_id] = _NodeDef(node=node, explicit=explicit)
+        return node
+
+    if existing.node == node:
+        return existing.node
+
+    if not explicit and existing.explicit:
+        return existing.node
+
+    if explicit and not existing.explicit:
+        nodes[spec.node_id] = _NodeDef(node=node, explicit=True)
+        return node
+
+    raise PromptFlowParseError(
+        _line_error(line_no, f'Conflicting definition for node "{spec.node_id}"')
+    )
+
+
+def _line_error(line_no: int, message: str) -> str:
+    return f"Line {line_no}: {message}"