diff options
Diffstat (limited to 'assistant.py')
| -rw-r--r-- | assistant.py | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/assistant.py b/assistant.py index 4ea8113..4dc11db 100644 --- a/assistant.py +++ b/assistant.py @@ -253,6 +253,7 @@ def llm_model(): def ask(messages, calls, tools=None): # Yields reply text as it streams; tool-call fragments go into calls. body = {"model": llm_model(), "messages": messages, "stream": True, + "stream_options": {"include_usage": True}, "chat_template_kwargs": {"enable_thinking": False}} if tools: body["tools"] = tools @@ -262,7 +263,12 @@ def ask(messages, calls, tools=None): for line in r: if not line.startswith(b"data: ") or line.strip() == b"data: [DONE]": continue - choices = json.loads(line[6:])["choices"] + chunk = json.loads(line[6:]) + if chunk.get("usage"): + # How full the model's context is: history, tool results and + # tool definitions all count. + log(f"context: {chunk['usage']['prompt_tokens']} prompt tokens") + choices = chunk["choices"] if not choices: continue delta = choices[0]["delta"] |