codegen output token stream prints to console in testing
This commit is contained in:
@@ -214,6 +214,11 @@ unit tests (24) + box integration tests (2).
|
||||
raised as `CodegenError` by the client, never a raw `TimeoutError`. The
|
||||
default codegen timeout is 1200s (`cli.build_scheduler`); raise
|
||||
`codegen.timeout` in config for harder prompts.
|
||||
- Set `codegen.stream: true` to echo the codegen output as an SSE token stream
|
||||
to stdout during body writes — including the chain-of-thought, so a long
|
||||
(~7 min) write shows live progress. Echoing is console-only; the returned
|
||||
content is identical either way. Integration tests already force streaming;
|
||||
see it with `-s` on the box.
|
||||
|
||||
### Code principles
|
||||
- **No mocking.** The decision engine is always real SemIf; the LLM is always a
|
||||
|
||||
@@ -53,6 +53,7 @@ def build_scheduler(config: dict) -> tuple[Scheduler, dict]:
|
||||
),
|
||||
model=codegen_cfg.get("model", "qwen38-iq3s"),
|
||||
timeout=float(codegen_cfg.get("timeout", 1200.0)),
|
||||
stream=bool(codegen_cfg.get("stream", False)),
|
||||
)
|
||||
log = DecisionLog(config.get("log", "data/decisions.jsonl"))
|
||||
trace = TraceLog(config.get("trace", "data/runs.jsonl"))
|
||||
|
||||
+52
-1
@@ -40,12 +40,24 @@ class CodegenClient:
|
||||
truncates the hidden reasoning, leaving `content` empty. Omit `max_tokens`
|
||||
so the model runs to completion; the reasoning is filtered automatically
|
||||
because only `content` is read.
|
||||
|
||||
With `stream=True` the response is read as an SSE token stream and echoed
|
||||
to stdout as it arrives — including the chain-of-thought — so a long
|
||||
body write shows live progress. Echoing is console-only; the returned
|
||||
content is identical either way.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str, model: str, timeout: float = 1200.0):
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
model: str,
|
||||
timeout: float = 1200.0,
|
||||
stream: bool = False,
|
||||
):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.stream = stream
|
||||
|
||||
def chat(
|
||||
self,
|
||||
@@ -58,6 +70,7 @@ class CodegenClient:
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"stream": self.stream,
|
||||
}
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
@@ -67,6 +80,8 @@ class CodegenClient:
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
||||
if self.stream:
|
||||
return self._read_stream(response)
|
||||
payload = json.loads(response.read().decode("utf-8"))
|
||||
except urllib.error.URLError as exc:
|
||||
raise CodegenError(
|
||||
@@ -78,6 +93,40 @@ class CodegenClient:
|
||||
) from exc
|
||||
return payload["choices"][0]["message"]["content"]
|
||||
|
||||
def _read_stream(self, response) -> str:
|
||||
"""Read an OpenAI-compatible SSE stream, echo tokens to stdout.
|
||||
|
||||
Only `content` deltas are accumulated into the returned body;
|
||||
`reasoning_content` (chain-of-thought) is echoed to the console but
|
||||
never part of the result.
|
||||
"""
|
||||
import sys
|
||||
|
||||
parts: list[str] = []
|
||||
for raw in response:
|
||||
line = raw.decode("utf-8").strip()
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data = line[len("data:") :].strip()
|
||||
if data == "[DONE]":
|
||||
break
|
||||
try:
|
||||
chunk = json.loads(data)
|
||||
except ValueError:
|
||||
continue
|
||||
choice = chunk.get("choices", [{}])[0]
|
||||
delta = choice.get("delta", {}) or {}
|
||||
text = delta.get("content") or ""
|
||||
reasoning = delta.get("reasoning_content") or ""
|
||||
if text or reasoning:
|
||||
sys.stdout.write(text + reasoning)
|
||||
sys.stdout.flush()
|
||||
parts.append(text)
|
||||
if parts:
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def build_skill_body_prompt(
|
||||
request: Request,
|
||||
@@ -165,6 +214,8 @@ def generate_skill_body(
|
||||
"""Author a skill body with the big model; retries once on invalid output."""
|
||||
contract_text = contract if contract is not None else read_skill_contract()
|
||||
messages = build_skill_body_prompt(request, category, draft, tree, contract_text)
|
||||
if client.stream:
|
||||
print(f"[codegen] writing body for {category}.{draft.name}...")
|
||||
last_error: Exception | None = None
|
||||
for attempt in range(2):
|
||||
try:
|
||||
|
||||
@@ -163,6 +163,7 @@ def test_generate_skill_body_codegen(tmp_path):
|
||||
base_url=codegen_cfg.get("base_url", "http://localhost:11434/v1"),
|
||||
model=codegen_cfg.get("model", "qwen38-iq3s"),
|
||||
timeout=float(codegen_cfg.get("timeout", 1200.0)),
|
||||
stream=True,
|
||||
)
|
||||
tree = build_tree(build_skills({"skills": {}}))
|
||||
draft = SkillDraft(
|
||||
@@ -223,6 +224,7 @@ def test_create_category_chain_runs_new_skill(tmp_path):
|
||||
config["trace"] = str(tmp_path / "runs.jsonl")
|
||||
config["category_registry"] = str(tmp_path / "categories.json")
|
||||
config["skill_bodies"] = str(tmp_path / "skills")
|
||||
config["codegen"] = {**config.get("codegen", {}), "stream": True}
|
||||
scheduler, config = build_scheduler(config)
|
||||
scheduler.tree = {}
|
||||
|
||||
|
||||
@@ -252,6 +252,79 @@ def test_chat_timeout_raises_codegen_error():
|
||||
httpd.server_close()
|
||||
|
||||
|
||||
def _sse_frame(payload: dict) -> str:
|
||||
return f"data: {json.dumps(payload)}\n\n"
|
||||
|
||||
|
||||
class _StreamingOpenAI(BaseHTTPRequestHandler):
|
||||
"""Replies with an OpenAI-compatible SSE token stream (COT then content)."""
|
||||
|
||||
reasoning: str = "thinking about the body..."
|
||||
content: str = GOOD_BODY
|
||||
received: list = []
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length") or 0)
|
||||
raw = self.rfile.read(length).decode("utf-8")
|
||||
type(self).received.append(json.loads(raw))
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream")
|
||||
self.end_headers()
|
||||
reasoning = type(self).reasoning
|
||||
content = type(self).content
|
||||
step = max(len(reasoning) // 4, 1)
|
||||
for i in range(0, len(reasoning), step):
|
||||
frame = _sse_frame(
|
||||
{"choices": [{"delta": {"reasoning_content": reasoning[i : i + step]}}]}
|
||||
)
|
||||
self.wfile.write(frame.encode("utf-8"))
|
||||
step = max(len(content) // 4, 1)
|
||||
for i in range(0, len(content), step):
|
||||
frame = _sse_frame(
|
||||
{"choices": [{"delta": {"content": content[i : i + step]}}]}
|
||||
)
|
||||
self.wfile.write(frame.encode("utf-8"))
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
self.close_connection = True
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
|
||||
def _streaming_server() -> tuple[ThreadingHTTPServer, str]:
|
||||
handler = type("Handler", (_StreamingOpenAI,), {"received": []})
|
||||
httpd = ThreadingHTTPServer(("127.0.0.1", 0), handler)
|
||||
threading.Thread(target=httpd.serve_forever, daemon=True).start()
|
||||
return httpd, f"http://127.0.0.1:{httpd.server_address[1]}/v1"
|
||||
|
||||
|
||||
def test_chat_stream_accumulates_full_content():
|
||||
httpd, base = _streaming_server()
|
||||
try:
|
||||
client = CodegenClient(base_url=base, model="test", timeout=10, stream=True)
|
||||
out = client.chat([{"role": "user", "content": "hi"}])
|
||||
assert out == GOOD_BODY, "streamed deltas must reassemble the full body"
|
||||
body = httpd.RequestHandlerClass.received[0]
|
||||
assert body["stream"] is True
|
||||
finally:
|
||||
httpd.shutdown()
|
||||
httpd.server_close()
|
||||
|
||||
|
||||
def test_chat_stream_verbose_echoes_tokens(capsys):
|
||||
httpd, base = _streaming_server()
|
||||
try:
|
||||
client = CodegenClient(base_url=base, model="test", timeout=10, stream=True)
|
||||
out = client.chat([{"role": "user", "content": "hi"}])
|
||||
assert out == GOOD_BODY
|
||||
captured = capsys.readouterr().out
|
||||
assert captured == _StreamingOpenAI.reasoning + GOOD_BODY + "\n"
|
||||
finally:
|
||||
httpd.shutdown()
|
||||
httpd.server_close()
|
||||
|
||||
|
||||
def test_chat_omits_max_tokens_by_default():
|
||||
httpd, base = _fake_server(GOOD_BODY)
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user