From f16bef46713e3f4266983586242090b74eb3c374 Mon Sep 17 00:00:00 2001 From: gramps Date: Mon, 13 Jul 2026 08:07:26 -0700 Subject: [PATCH] tok: #/## % badge with context tracking, client-side token counting, remove triage dep --- config.py | 1 + routers/chat.py | 13 ++++++++----- templates/index.html | 9 ++++++++- tests/test_chat_streaming_and_memory_paths.py | 4 +--- 4 files changed, 18 insertions(+), 9 deletions(-) diff --git a/config.py b/config.py index 3bf87cd..522741d 100644 --- a/config.py +++ b/config.py @@ -15,6 +15,7 @@ LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8 SEARXNG_BASE = "http://localhost:8888" DEFAULT_MODEL = "llama3.1:latest" COMPLETIONS_API_KEY = os.environ.get("CAIC_COMPLETIONS_API_KEY", "caic-sk-" + os.urandom(24).hex()) +MODEL_CONTEXT_LENGTH = 4096 # --- AMQP --- AMQP_RECONNECT_DELAY = 5 diff --git a/routers/chat.py b/routers/chat.py index ebd8ac6..0fc9232 100644 --- a/routers/chat.py +++ b/routers/chat.py @@ -16,7 +16,7 @@ from search import (calculate_perplexity, is_uncertain, is_refusal, clean_hedging, format_search_results, format_direct_answer, extract_search_query, query_searxng) from security import read_json_body, log_incident, BODY_LIMIT_CHAT_BYTES -from config import MAX_CHAT_MESSAGE_CHARS +from config import MAX_CHAT_MESSAGE_CHARS, MODEL_CONTEXT_LENGTH log = logging.getLogger("caic") router = APIRouter() @@ -46,6 +46,7 @@ def parse_llama_stream_chunk(line: str) -> tuple: usage = chunk.get("usage", {}) stats["tokens_per_sec"] = usage.get("tokens_per_second", 0.0) stats["completion_tokens"] = usage.get("completion_tokens", 0) + stats["prompt_tokens"] = usage.get("prompt_tokens", 0) return token, finish == "stop", stats, logprobs_list if "message" in chunk and "content" in chunk["message"]: token = chunk["message"]["content"] @@ -125,6 +126,7 @@ async def chat(request: Request): all_logprobs = [] tokens_per_sec = 0.0 completion_tokens = 0 + prompt_tokens = 0 if remember_response: yield f"data: {json.dumps({'token': remember_response + chr(10) + chr(10), 'conversation_id': conv_id})}\n\n" @@ -132,7 +134,7 @@ async def chat(request: Request): async with httpx.AsyncClient() as client: try: async with client.stream( - "POST", f"{inference_base}/v1/chat/completions", + "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions", json=upstream_payload, timeout=httpx.Timeout(300.0, connect=10.0), ) as resp: @@ -147,6 +149,7 @@ async def chat(request: Request): if done: tokens_per_sec = stats.get("tokens_per_sec", 0.0) completion_tokens = stats.get("completion_tokens", 0) + prompt_tokens = stats.get("prompt_tokens", 0) assistant_msg = "".join(full_response) perplexity = calculate_perplexity(all_logprobs) if all_logprobs else 0.0 @@ -172,7 +175,7 @@ async def chat(request: Request): augmented_response = [] async with client.stream( - "POST", f"{inference_base}/v1/chat/completions", + "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions", json={"model": model, "messages": augmented_messages, "stream": True}, timeout=httpx.Timeout(300.0, connect=10.0), ) as resp2: @@ -201,7 +204,7 @@ async def chat(request: Request): db2.commit() db2.close() - yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'searched': True, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'tokens': completion_tokens})}\n\n" + yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'searched': True, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH})}\n\n" return saved_msg = assistant_msg @@ -214,7 +217,7 @@ async def chat(request: Request): db2.commit() db2.close() - yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'tokens': completion_tokens})}\n\n" + yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH})}\n\n" except httpx.RemoteProtocolError: pass diff --git a/templates/index.html b/templates/index.html index a03815d..29ffe59 100644 --- a/templates/index.html +++ b/templates/index.html @@ -1405,7 +1405,14 @@ async function sendMessage() { const ttrClass = data.searched ? 'search' : (ttrSec <= 2 ? 'low' : ttrSec < 11 ? 'medium' : 'high'); roleLabel.innerHTML += `ttr: ${ttrSec.toFixed(1)}s`; } - if (tokenCount > 0 && roleLabel) roleLabel.innerHTML += `tok: ${tokenCount}`; + if (tokenCount > 0 && roleLabel) { + const pt = data.prompt_tokens || 0; + const ct = data.completion_tokens || 0; + const totalTokens = pt + ct; + const ctxLen = data.context_length || 4096; + const usedPct = Math.round((totalTokens / ctxLen) * 100); + roleLabel.innerHTML += `tok: ${totalTokens}/${ctxLen} ${usedPct}%`; + } conversationHistory.push({ role: 'assistant', content: fullText }); updateTokenThermometer(); addCopyButtons(assistantDiv); diff --git a/tests/test_chat_streaming_and_memory_paths.py b/tests/test_chat_streaming_and_memory_paths.py index d8cfd1b..23b3243 100644 --- a/tests/test_chat_streaming_and_memory_paths.py +++ b/tests/test_chat_streaming_and_memory_paths.py @@ -14,9 +14,7 @@ from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS def _mock_triage_url(monkeypatch, url: str = config.LLAMA_SERVER_BASE): - async def fake_url(query: str) -> str: - return url - monkeypatch.setattr(routers.chat, "_get_inference_url", fake_url) + monkeypatch.setattr(config, "LLAMA_SERVER_BASE", url) def make_client(tmp_path: Path) -> TestClient: