| import json |
| import time |
| from collections.abc import Iterator |
| import httpx |
|
|
| from open_cortex.runtime.events import RuntimeEvent |
| from open_cortex.runtime.metrics import fetch_runtime_snapshot |
| from open_cortex.runtime.messages import ChatMessage, to_llama_messages |
|
|
|
|
| CHAT_URL = "http://127.0.0.1:8080/v1/chat/completions" |
| METRICS_URL = "http://127.0.0.1:8080/metrics" |
| SLOTS_URL = "http://127.0.0.1:8080/slots" |
|
|
|
|
| def _raise_for_status_with_body(response: httpx.Response) -> None: |
| if response.is_error: |
| try: |
| response.read() |
| except httpx.HTTPError: |
| pass |
| response.raise_for_status() |
|
|
|
|
| def _working_memory_percent( |
| context_tokens: int | None, |
| context_size: int | None, |
| ) -> float | None: |
| if context_tokens is None or not context_size: |
| return None |
| return round(min(100.0, context_tokens / context_size * 100), 1) |
|
|
|
|
| def _detect_repetition(text: str) -> bool: |
| normalized = " ".join(text.split()) |
| if len(normalized) < 120: |
| return False |
|
|
| for window in (96, 72, 48, 32): |
| if len(normalized) <= window * 2: |
| continue |
| tail = normalized[-window:] |
| if normalized[:-window].count(tail) >= 1: |
| return True |
|
|
| return False |
|
|
|
|
| def stream_chat_events(message: list[ChatMessage]) -> Iterator[RuntimeEvent]: |
| request_body = { |
| "messages": to_llama_messages(message), |
| "temperature": 0.2, |
| "max_tokens": 1024, |
| "stream": True, |
| "stream_options": {"include_usage": True}, |
| "timings_per_token": True, |
| } |
|
|
| request_started = time.perf_counter() |
| first_token_seen = False |
| first_token_at = None |
| generated_tokens = 0 |
| generated_text = "" |
| final_stats = None |
| base_context_tokens = None |
| context_size = None |
|
|
| yield RuntimeEvent( |
| kind="request_started", |
| text_delta="", |
| ttft_ms=None, |
| snapshot=None, |
| ) |
|
|
| with httpx.Client(timeout=120.0, trust_env=False) as client: |
| with client.stream("POST", CHAT_URL, json=request_body) as response: |
| _raise_for_status_with_body(response) |
|
|
| for line in response.iter_lines(): |
| if not line.startswith("data: "): |
| continue |
|
|
| data = line.removeprefix("data: ") |
|
|
| if data == "[DONE]": |
| break |
|
|
| event = json.loads(data) |
| choices = event.get("choices", []) |
|
|
| if choices: |
| content = choices[0].get("delta", {}).get("content") |
|
|
| if content: |
| now = time.perf_counter() |
| generated_tokens += 1 |
| generated_text += content |
| elapsed_ms = ( |
| (now - first_token_at) * 1000 |
| if first_token_at is not None |
| else 0.0 |
| ) |
| live_tps = ( |
| generated_tokens / (elapsed_ms / 1000) |
| if elapsed_ms > 0 |
| else None |
| ) |
| repetition_detected = _detect_repetition(generated_text) |
|
|
| if not first_token_seen: |
| first_token_seen = True |
| first_token_at = now |
| ttft_ms = (first_token_at - request_started) * 1000 |
| snapshot = fetch_runtime_snapshot( |
| client, |
| METRICS_URL, |
| SLOTS_URL, |
| ) |
| base_context_tokens = ( |
| snapshot.slot_context_tokens[0] |
| if snapshot.slot_context_tokens |
| else None |
| ) |
| context_size = snapshot.slot_context_size |
| context_tokens = ( |
| base_context_tokens + generated_tokens |
| if base_context_tokens is not None |
| else None |
| ) |
| yield RuntimeEvent( |
| kind="first_token", |
| text_delta=content, |
| ttft_ms=ttft_ms, |
| snapshot=snapshot, |
| generated_tokens=generated_tokens, |
| elapsed_ms=0.0, |
| live_tps=None, |
| repetition_detected=repetition_detected, |
| context_tokens=context_tokens, |
| context_size=context_size, |
| working_memory_percent=_working_memory_percent( |
| context_tokens, |
| context_size, |
| ), |
| ) |
| else: |
| context_tokens = ( |
| base_context_tokens + generated_tokens |
| if base_context_tokens is not None |
| else None |
| ) |
| yield RuntimeEvent( |
| kind="token", |
| text_delta=content, |
| ttft_ms=None, |
| snapshot=None, |
| generated_tokens=generated_tokens, |
| elapsed_ms=elapsed_ms, |
| live_tps=live_tps, |
| repetition_detected=repetition_detected, |
| context_tokens=context_tokens, |
| context_size=context_size, |
| working_memory_percent=_working_memory_percent( |
| context_tokens, |
| context_size, |
| ), |
| ) |
|
|
| if event.get("usage"): |
| final_stats = event |
|
|
| if final_stats is not None: |
| usage = final_stats["usage"] |
| timings = final_stats["timings"] |
| context_tokens = usage["prompt_tokens"] + usage["completion_tokens"] |
|
|
| yield RuntimeEvent( |
| kind="request_completed", |
| text_delta="", |
| ttft_ms=None, |
| snapshot=None, |
| prompt_tokens=usage["prompt_tokens"], |
| completion_tokens=usage["completion_tokens"], |
| prompt_tps=timings["prompt_per_second"], |
| decode_tps=timings["predicted_per_second"], |
| generated_tokens=usage["completion_tokens"], |
| repetition_detected=_detect_repetition(generated_text), |
| context_tokens=context_tokens, |
| context_size=context_size, |
| working_memory_percent=_working_memory_percent( |
| context_tokens, |
| context_size, |
| ), |
| ) |
|
|
