mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 06:15:58 -10:00
chat.py: fenced code blocks in streamed answers are coloured on a terminal; prose is untouched (PR #1895 by MegaRushX)
This commit is contained in:
MegaRushX authored and Niko1221 committed
1 parent
df6928400f
commit
3e81173d63
3 files changed
+223
-1
No files matched your search
@@ -18,6 +18,8 @@ import sys
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
from highlight import StreamHighlighter
|
||||
|
||||
|
||||
def stream(url, messages, think, max_tokens):
|
||||
body = {"model": "strata", "messages": messages, "stream": True, "max_tokens": max_tokens,
|
||||
@@ -43,6 +45,7 @@ def main() -> int:
|
||||
a = ap.parse_args()
|
||||
url = f"http://{a.host}:{a.port}/v1/chat/completions"
|
||||
gray, reset = ("\033[90m", "\033[0m") if sys.stdout.isatty() else ("", "")
|
||||
hl = StreamHighlighter(color=sys.stdout.isatty())
|
||||
messages, pending = [], []
|
||||
think = "none" if a.no_think else a.think
|
||||
print(f"Strata chat ({url}). /image <path> = attach a picture, /think none|low|medium|high (now: {think}), "
|
||||
@@ -98,11 +101,14 @@ def main() -> int:
|
||||
if in_think:
|
||||
print(reset + "\n", end="")
|
||||
in_think = False
|
||||
print(content, end="", flush=True)
|
||||
print(hl.feed(content), end="", flush=True)
|
||||
answer.append(content)
|
||||
n += 1
|
||||
if in_think:
|
||||
print(reset, end="")
|
||||
out = hl.flush()
|
||||
if out:
|
||||
print(out, end="")
|
||||
except OSError as e:
|
||||
print(f"\n(could not reach the server at {url}: {e})")
|
||||
messages.pop()
|
||||
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
"""Syntax highlighting for chat.py's streamed output: fenced code blocks get colored (tags, strings, comments,
|
||||
keywords, numbers), everything else passes through unchanged. Standard library only, no dependency.
|
||||
|
||||
from highlight import StreamHighlighter
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed(streamed_chunk) # complete lines, highlighted; a partial last line is held back
|
||||
out += hl.flush() # the held-back tail (unclosed fence included)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
|
||||
# ANSI colors; a disabled colorer returns the text unchanged
|
||||
class _Palette:
|
||||
def __init__(self, color: bool):
|
||||
self.on = color
|
||||
|
||||
def paint(self, text: str, code: str) -> str:
|
||||
return f"\033[{code}m{text}\033[0m" if self.on else text
|
||||
|
||||
|
||||
# one regex per language family; the first matching family paints the line
|
||||
_PY_KEYWORDS = ("False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue",
|
||||
"def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import", "in",
|
||||
"is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield")
|
||||
_JS_KEYWORDS = ("async", "await", "break", "case", "catch", "class", "const", "continue", "default", "delete",
|
||||
"do", "else", "export", "extends", "false", "finally", "for", "function", "if", "import", "in",
|
||||
"instanceof", "let", "new", "null", "return", "static", "super", "switch", "this", "throw", "true",
|
||||
"try", "typeof", "undefined", "var", "void", "while", "yield")
|
||||
_C_KEYWORDS = ("alignas", "auto", "bool", "break", "case", "catch", "char", "class", "const", "constexpr",
|
||||
"continue", "default", "delete", "do", "double", "else", "enum", "explicit", "extern", "false",
|
||||
"float", "for", "friend", "if", "inline", "int", "long", "namespace", "new", "noexcept", "nullptr",
|
||||
"private", "protected", "public", "return", "short", "signed", "sizeof", "static", "struct",
|
||||
"switch", "template", "this", "throw", "true", "try", "typedef", "typename", "union", "unsigned",
|
||||
"using", "virtual", "void", "while")
|
||||
_SHELL_KEYWORDS = ("case", "do", "done", "elif", "else", "esac", "fi", "for", "function", "if", "in", "then",
|
||||
"until", "while")
|
||||
_JSON_KEYS = ("false", "null", "true")
|
||||
|
||||
_STR = r"(\"(?:[^\"\\]|\\.)*\")" # "..." with escapes
|
||||
_SQ_STR = r"('(?:[^'\\]|\\.)*')"
|
||||
_COMMENT = r"(#.*$)"
|
||||
_SLASH_COMMENT = r"(//.*$)"
|
||||
_NUM = r"\b(\d[\w.]*)\b"
|
||||
|
||||
_FAMILIES = {
|
||||
"markup": ({"html", "htm", "xml", "svg", "vue", "jsx", "tsx"},
|
||||
re.compile(r"(<!--.*?-->)|(<[A-Za-z][\w:.-]*/?>?)|(</[A-Za-z][\w:.-]*>)|([A-Za-z-]+)(?==)"),
|
||||
{1: "36;1", 2: "96", 3: "96", 4: "94"}),
|
||||
"python": ({"py", "python"},
|
||||
re.compile(rf"{_COMMENT}|{_SQ_STR}|{_STR}|\b({'|'.join(_PY_KEYWORDS)})\b|{_NUM}"),
|
||||
{1: "90", 2: "33", 3: "33", 4: "35;1", 5: "34"}),
|
||||
"c-like": ({"c", "cc", "cpp", "c++", "h", "hpp", "cu", "hip", "java", "kt", "rs", "go", "cs"},
|
||||
re.compile(rf"{_SLASH_COMMENT}|{_STR}|\b({'|'.join(_C_KEYWORDS)})\b|\b(0x[0-9a-fA-F]+|\d[\w.]*)\b"),
|
||||
{1: "90", 2: "33", 3: "35;1", 4: "34"}),
|
||||
"js-like": ({"js", "javascript", "ts", "typescript", "mjs"},
|
||||
re.compile(rf"{_SLASH_COMMENT}|{_STR}|{_SQ_STR}|(`[^`]*`)|\b({'|'.join(_JS_KEYWORDS)})\b|{_NUM}"),
|
||||
{1: "90", 2: "33", 3: "33", 4: "33", 5: "35;1", 6: "34"}),
|
||||
"shell": ({"sh", "bash", "zsh", "shell", "console", "bat", "cmd", "powershell", "ps1"},
|
||||
re.compile(rf"{_COMMENT}|{_SQ_STR}|{_STR}|\b({'|'.join(_SHELL_KEYWORDS)})\b|(^\s*\$)"),
|
||||
{1: "90", 2: "33", 3: "33", 4: "35;1", 5: "90"}),
|
||||
"json": ({"json", "jsonc", "json5"},
|
||||
re.compile(rf"{_SLASH_COMMENT}|{_STR}(?=\s*:)|\b({'|'.join(_JSON_KEYS)})\b|\b(-?\d[\w.]*)\b"),
|
||||
{1: "90", 2: "94", 3: "35;1", 4: "34"}),
|
||||
"css": ({"css", "scss", "less"},
|
||||
re.compile(r"(/\*.*?\*/)|([.#][\w-]+)|([\w-]+)(?=\s*:)|\b(\d[\w.%]*)\b"),
|
||||
{1: "90", 2: "96", 3: "94", 4: "34"}),
|
||||
"sql": ({"sql"},
|
||||
re.compile(r"(--.*$)|('[^']*')|\b(select|from|where|insert|into|update|set|delete|join|left|right|inner|"
|
||||
r"outer|on|group|by|order|having|limit|create|table|alter|drop|and|or|not|null|as)\b",
|
||||
re.IGNORECASE),
|
||||
{1: "90", 2: "33", 3: "35;1"}),
|
||||
"diff": ({"diff", "patch"},
|
||||
re.compile(r"^(\+[^+]*)|^(-.*)|^(@@.*@@)"),
|
||||
{1: "32", 2: "31", 3: "36"}),
|
||||
}
|
||||
|
||||
_FENCE_RE = re.compile(r"^\s*(```|~~~)\s*([\w+#-]*)")
|
||||
_LANG_ALIAS = {"c++": "c-like", "javascript": "js-like", "typescript": "js-like", "python": "python",
|
||||
"py": "python", "shell": "shell", "bash": "shell", "console": "shell"}
|
||||
|
||||
|
||||
def _family(lang: str):
|
||||
lang = _LANG_ALIAS.get(lang.lower(), lang.lower())
|
||||
for exts, rx, colors in _FAMILIES.values():
|
||||
if lang in exts:
|
||||
return rx, colors
|
||||
return None, None
|
||||
|
||||
|
||||
class StreamHighlighter:
|
||||
"""Line-buffered highlighter for streamed text. feed() returns complete lines (highlighted inside fenced code
|
||||
blocks); flush() returns the held-back tail. Outside code fences the text is returned unchanged (color on or
|
||||
off), so prose never gets painted. All characters other than the added ANSI escapes pass through unchanged."""
|
||||
|
||||
def __init__(self, color: bool = True):
|
||||
self.color = color
|
||||
self._pal = _Palette(color)
|
||||
self._buf = "" # held-back partial line
|
||||
self._in_fence = False
|
||||
self._rx = None
|
||||
self._colors = None
|
||||
|
||||
def feed(self, chunk: str) -> str:
|
||||
self._buf += chunk
|
||||
out = []
|
||||
while "\n" in self._buf:
|
||||
line, self._buf = self._buf.split("\n", 1)
|
||||
out.append(self._line(line))
|
||||
out.append("\n")
|
||||
return "".join(out)
|
||||
|
||||
def flush(self) -> str:
|
||||
tail, self._buf = self._buf, ""
|
||||
return self._line(tail) if tail else ""
|
||||
|
||||
def _line(self, line: str) -> str:
|
||||
fence = _FENCE_RE.match(line)
|
||||
if fence:
|
||||
marker, lang = fence.group(1), fence.group(2)
|
||||
if self._in_fence:
|
||||
self._in_fence, self._rx, self._colors = False, None, None
|
||||
return self._pal.paint(line, "90")
|
||||
self._in_fence = True
|
||||
self._rx, self._colors = _family(lang)
|
||||
if not self._rx:
|
||||
return self._pal.paint(line, "90")
|
||||
return self._pal.paint(marker, "90") + self._pal.paint(lang, "94") + line[fence.end():]
|
||||
if not self._in_fence or not self._rx:
|
||||
return line
|
||||
return self._rx.sub(self._sub, line)
|
||||
|
||||
def _sub(self, m: re.Match) -> str:
|
||||
for gi, code in self._colors.items():
|
||||
if m.group(gi) is not None:
|
||||
return self._pal.paint(m.group(gi), code)
|
||||
return m.group(0)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""highlight.py: fenced code blocks in chat.py's streamed output get syntax highlighting (tags, strings, comments,
|
||||
keywords, numbers), prose passes through unchanged, and a fence split across streamed chunks still opens.
|
||||
|
||||
python -m pytest tests/terminal/test_chat_highlight.py (no GPU, no downloads)
|
||||
"""
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from highlight import StreamHighlighter # noqa: E402
|
||||
|
||||
RESET = "\033[0m"
|
||||
|
||||
|
||||
def colored(text, code):
|
||||
return f"\033[{code}m{text}{RESET}"
|
||||
|
||||
|
||||
class TestHighlight(unittest.TestCase):
|
||||
def test_prose_unchanged_with_color(self):
|
||||
hl = StreamHighlighter(color=True)
|
||||
self.assertEqual(hl.feed("hello world\n"), "hello world\n")
|
||||
self.assertEqual(hl.flush(), "")
|
||||
|
||||
def test_prose_unchanged_without_color(self):
|
||||
hl = StreamHighlighter(color=False)
|
||||
self.assertEqual(hl.feed("def f(x):\n return x\n"), "def f(x):\n return x\n")
|
||||
|
||||
def test_python_block_is_highlighted(self):
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed("```python\nfor i in range(3): # loop\n print(i)\n```\n")
|
||||
self.assertIn(colored("for", "35;1"), out)
|
||||
self.assertIn(colored("# loop", "90"), out)
|
||||
self.assertIn(colored("in", "35;1"), out)
|
||||
lines = out.split("\n")
|
||||
self.assertTrue(lines[-1] == "" and lines[-2] == colored("```", "90"))
|
||||
# prose after the fence is unchanged
|
||||
self.assertEqual(hl.feed("plain again\n"), "plain again\n")
|
||||
|
||||
def test_html_tags_highlighted(self):
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed("```html\n<div class=\"box\">hi</div>\n```\n")
|
||||
self.assertIn(colored("<div", "96"), out)
|
||||
self.assertIn(colored("</div>", "96"), out)
|
||||
self.assertIn(colored("class", "94"), out)
|
||||
|
||||
def test_fence_split_across_chunks(self):
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed("code:\n``") + hl.feed("`py\nimport os\n```\n")
|
||||
self.assertIn(colored("import", "35;1"), out)
|
||||
self.assertIn(colored("```", "90"), out)
|
||||
self.assertIn(colored("py", "94"), out)
|
||||
|
||||
def test_partial_line_held_back_then_flushed(self):
|
||||
hl = StreamHighlighter(color=False)
|
||||
self.assertEqual(hl.feed("abc"), "")
|
||||
self.assertEqual(hl.feed("def\nxyz"), "abcdef\n")
|
||||
self.assertEqual(hl.flush(), "xyz")
|
||||
|
||||
def test_unknown_language_passes_through(self):
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed("```brainfuck\n+++++\n```\n")
|
||||
self.assertIn("+++++", out)
|
||||
self.assertIn(colored("```brainfuck", "90"), out)
|
||||
|
||||
def test_streaming_bytes_preserved(self):
|
||||
text = "```python\ndef f(x):\n return x**2\n```\ndone\n"
|
||||
hl = StreamHighlighter(color=True)
|
||||
out = hl.feed(text)
|
||||
self.assertEqual(out.replace(RESET, "").replace("\033[35;1m", "").replace("\033[90m", "")
|
||||
.replace("\033[33m", "").replace("\033[34m", "").replace("\033[94m", ""), text)
|
||||
self.assertEqual(hl.flush(), "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user