Add cost tracking: extract usage from streaming SSE responses
ober
e888683b5c6d8696940dfb55ac77fe0547600235
--- a/lib/jcode/core/agent.sls +++ b/lib/jcode/core/agent.sls @@ -4,7 +4,7 @@ (library (jcode core agent) (export agent-run agent-chat agent-step current-stream-cb - current-tool-cb current-provider-override + current-tool-cb current-usage-cb current-provider-override current-model-override) (import (except (chezscheme) make-hash-table hash-table? iota \x31;+ \x31;- @@ -23,6 +23,7 @@ (current-directory))) (def current-stream-cb (make-parameter #f)) (def current-tool-cb (make-parameter #f)) + (def current-usage-cb (make-parameter #f)) (def (agent-run session-id user-input) (log-info logger "agent-run" `((session . ,session-id))) (let ([existing (session-get-messages session-id)]) @@ -58,12 +59,14 @@ (def (agent-loop-stream session-id messages) (let* ([provider (get-current-provider)] [tools (get-tool-schemas)]) - (let-values ([(content tool-calls) + (let-values ([(content tool-calls usage) (provider-stream-chat provider messages tools (current-stream-cb))]) + (when (and usage (current-usage-cb)) + ((current-usage-cb) usage)) (let ([response (make-assistant-message (if (string=? content "") #f content) (if (null? tool-calls) #f tool-calls))]) @@ -124,7 +127,7 @@ (agent-chat-loop provider new-messages tools)) (message-content response)))) (def (agent-chat-loop-stream provider messages tools) - (let-values ([(content tool-calls) + (let-values ([(content tool-calls usage) (provider-stream-chat provider messages --- a/lib/jcode/provider/provider.sls +++ b/lib/jcode/provider/provider.sls @@ -359,6 +359,9 @@ (let ([body (make-hash-table)]) (hash-put! body "model" (provider-model provider)) (hash-put! body "stream" #t) + (let ([opts (make-hash-table)]) + (hash-put! opts "include_usage" #t) + (hash-put! body "stream_options" opts)) (hash-put! body "messages" (map message->json messages)) (when (and tools (not (null? tools))) (hash-put! body "tools" tools)) @@ -370,7 +373,8 @@ [headers (openai-headers provider)] [body (openai-stream-body provider messages tools)] [text-acc (open-output-string)] - [tc-table (make-hash-table)]) + [tc-table (make-hash-table)] + [usage-acc (make-hash-table)]) (let* ([body-json (json-object->string body)] [dummy (log-info logger @@ -408,6 +412,19 @@ (string->json-object data))]) (when json + (let ([usage (hash-get + json + "usage")]) + (when (and usage + (hash-table? + usage)) + (hash-for-each + (lambda (k v) + (hash-put! + usage-acc + k + v)) + usage))) (let* ([choices (hash-get json "choices")] @@ -502,8 +519,23 @@ logger "stream-result" `((content-len . ,(string-length content)) - (tool-calls . ,(length tool-calls)))) - (values content tool-calls)))) + (tool-calls . ,(length tool-calls)) + (tokens-in . ,(or (hash-get usage-acc "prompt_tokens") 0)) + (tokens-out + . + ,(or (hash-get usage-acc "completion_tokens") 0)) + (cost . ,(or (hash-get usage-acc "cost") 0)))) + (values + content + tool-calls + (list + (cons + 'tokens-in + (or (hash-get usage-acc "prompt_tokens") 0)) + (cons + 'tokens-out + (or (hash-get usage-acc "completion_tokens") 0)) + (cons 'cost (or (hash-get usage-acc "cost") 0))))))) (def (anthropic-stream-headers provider) `(("Content-Type" . "application/json") ("x-api-key" . ,(provider-api-key provider)) @@ -524,7 +556,8 @@ [body (anthropic-stream-body provider messages tools)] [text-acc (open-output-string)] [tu-table (make-hash-table)] - [current-idx (make-parameter #f)]) + [current-idx (make-parameter #f)] + [usage-acc (make-hash-table)]) (let ([http-status (http-post-stream url headers @@ -641,6 +674,42 @@ event-type "content_block_stop") (current-idx #f)) + ((equal? + event-type + "message_start") + (let ([msg (hash-get + json + "message")]) + (when msg + (let ([usage (hash-get + msg + "usage")]) + (when (and usage + (hash-table? + usage)) + (hash-for-each + (lambda (k v) + (hash-put! + usage-acc + k + v)) + usage)))))) + ((equal? + event-type + "message_delta") + (let ([usage (hash-get + json + "usage")]) + (when (and usage + (hash-table? + usage)) + (hash-for-each + (lambda (k v) + (hash-put! + usage-acc + k + v)) + usage)))) (#t (void))]))))))))]) (unless (= http-status 200) (log-error @@ -657,7 +726,15 @@ (or (hash-get acc "name") "unknown") (or (hash-get acc "args") "{}")))) indices)]) - (values content tool-calls)))) + (values + content + tool-calls + (list + (cons 'tokens-in (or (hash-get usage-acc "input_tokens") 0)) + (cons + 'tokens-out + (or (hash-get usage-acc "output_tokens") 0)) + (cons 'cost 0)))))) (def (provider-stream-chat provider messages tools token-cb) (unless (provider-api-key provider) (error 'provider-stream-chat @@ -680,7 +757,7 @@ [content (or (message-content response) "")] [tcs (or (message-tool-calls response) '())]) (when (> (string-length content) 0) (token-cb content)) - (values content tcs))] + (values content tcs '()))] [else (error 'provider-stream-chat (format --- a/lib/jcode/ui/tui-input.sls +++ b/lib/jcode/ui/tui-input.sls @@ -48,11 +48,20 @@ (= key TB_KEY_CTRL_J)) 'submit] [(= key TB_KEY_ENTER) - (if (input-state-completion inp) - (begin (accept-completion! inp) 'continue) - (if (string-contains (input-state-text inp) "\n") - (begin (insert-char! inp #\newline) 'continue) - 'submit))] + (cond + [(input-state-completion inp) + (let* ([comp (input-state-completion inp)] + [selected (list-ref + comp + (input-state-comp-selected inp))] + [text (string-trim (input-state-text inp))]) + (if (equal? text (car selected)) + (begin (input-state-completion-set! inp #f) 'submit) + (begin (accept-completion! inp) 'continue)))] + [(string-contains (input-state-text inp) "\n") + (insert-char! inp #\newline) + 'continue] + [#t 'submit])] [(= key TB_KEY_CTRL_C) 'cancel] [(= key TB_KEY_CTRL_D) (if (string-empty? (input-state-text inp)) 'quit 'continue)] @@ -206,9 +215,7 @@ (let ([selected (list-ref comp (input-state-comp-selected inp))]) - (input-state-text-set! - inp - (string-append (car selected) " ")) + (input-state-text-set! inp (car selected)) (input-state-cursor-pos-set! inp (string-length (input-state-text inp))) --- a/lib/jcode/ui/tui.sls +++ b/lib/jcode/ui/tui.sls @@ -494,7 +494,10 @@ state event name - args))]) + args))] + [current-usage-cb + (lambda (usage) + (tui-usage-update! state usage))]) (agent-run s-id text) (tui-log "run-agent: agent-run returned, buf-len=~a" @@ -542,6 +545,25 @@ (reflow-message! m (msg-area-width state))) (loop (cdr msgs))))))) (app-state-dirty?-set! state #t)])) + (def (tui-usage-update! state usage) + "Accumulate token/cost usage from a streaming response." + (for-each + (lambda (pair) + (case (car pair) + [(tokens-in) + (app-state-tokens-in-set! + state + (+ (app-state-tokens-in state) (cdr pair)))] + [(tokens-out) + (app-state-tokens-out-set! + state + (+ (app-state-tokens-out state) (cdr pair)))] + [(cost) + (app-state-cost-set! + state + (+ (app-state-cost state) (cdr pair)))])) + usage) + (app-state-dirty?-set! state #t)) (def (tool-event-metadata name args) (cond [(hash-table? args) --- a/src/jcode/core/agent.ss +++ b/src/jcode/core/agent.ss @@ -5,6 +5,7 @@ agent-step current-stream-cb current-tool-cb + current-usage-cb current-provider-override current-model-override) @@ -43,6 +44,7 @@ Prefer using the edit tool over write for modifying existing files." (current-di (def current-stream-cb (make-parameter #f)) (def current-tool-cb (make-parameter #f)) +(def current-usage-cb (make-parameter #f)) (def (agent-run session-id user-input) (log-info logger "agent-run" `((session . ,session-id))) @@ -72,8 +74,10 @@ Prefer using the edit tool over write for modifying existing files." (current-di ;; Streaming version: calls (current-stream-cb) for each text token. (let* ((provider (get-current-provider)) (tools (get-tool-schemas))) - (let-values (((content tool-calls) + (let-values (((content tool-calls usage) (provider-stream-chat provider messages tools (current-stream-cb)))) + (when (and usage (current-usage-cb)) + ((current-usage-cb) usage)) (let ((response (make-assistant-message (if (string=? content "") #f content) (if (null? tool-calls) #f tool-calls)))) @@ -125,7 +129,7 @@ Prefer using the edit tool over write for modifying existing files." (current-di (message-content response)))) (def (agent-chat-loop-stream provider messages tools) - (let-values (((content tool-calls) + (let-values (((content tool-calls usage) (provider-stream-chat provider messages tools (current-stream-cb)))) (if (null? tool-calls) content --- a/src/jcode/provider/provider.ss +++ b/src/jcode/provider/provider.ss @@ -329,6 +329,10 @@ (let ((body (make-hash-table))) (hash-put! body "model" (provider-model provider)) (hash-put! body "stream" #t) + ;; Request usage data in stream + (let ((opts (make-hash-table))) + (hash-put! opts "include_usage" #t) + (hash-put! body "stream_options" opts)) (hash-put! body "messages" (map message->json messages)) (when (and tools (not (null? tools))) (hash-put! body "tools" tools)) @@ -336,13 +340,14 @@ (def (openai-stream-chat provider messages tools token-cb) ;; Stream via SSE. Calls token-cb with each text token. - ;; Returns (values content-string tool-call-list) + ;; Returns (values content-string tool-call-list usage-alist) (let* ((url (string-append (provider-base-url provider) "/chat/completions")) (headers (openai-headers provider)) (body (openai-stream-body provider messages tools)) (text-acc (open-output-string)) ;; tool-call accumulators: index -> alist with id/name/args-so-far - (tc-table (make-hash-table))) + (tc-table (make-hash-table)) + (usage-acc (make-hash-table))) (let* ((body-json (json-object->string body)) (dummy (log-info logger "stream-request" `((url . ,url) (body-len . ,(string-length body-json))))) @@ -363,6 +368,10 @@ (let ((json (guard (e [#t #f]) (string->json-object data)))) (when json + ;; Capture usage from final chunk + (let ((usage (hash-get json "usage"))) + (when (and usage (hash-table? usage)) + (hash-for-each (lambda (k v) (hash-put! usage-acc k v)) usage))) (let* ((choices (hash-get json "choices")) (choice (and (pair? choices) (car choices))) (delta (and choice (hash-get choice "delta")))) @@ -411,8 +420,14 @@ indices))) (log-info logger "stream-result" `((content-len . ,(string-length content)) - (tool-calls . ,(length tool-calls)))) - (values content tool-calls)))) + (tool-calls . ,(length tool-calls)) + (tokens-in . ,(or (hash-get usage-acc "prompt_tokens") 0)) + (tokens-out . ,(or (hash-get usage-acc "completion_tokens") 0)) + (cost . ,(or (hash-get usage-acc "cost") 0)))) + (values content tool-calls + (list (cons 'tokens-in (or (hash-get usage-acc "prompt_tokens") 0)) + (cons 'tokens-out (or (hash-get usage-acc "completion_tokens") 0)) + (cons 'cost (or (hash-get usage-acc "cost") 0))))))) ;;; Anthropic Streaming ;;; @@ -427,14 +442,15 @@ body)) (def (anthropic-stream-chat provider messages tools token-cb) - ;; Anthropic SSE streaming. Returns (values content tool-call-list) + ;; Anthropic SSE streaming. Returns (values content tool-call-list usage-alist) (let* ((url (string-append (provider-base-url provider) "/messages")) (headers (anthropic-stream-headers provider)) (body (anthropic-stream-body provider messages tools)) (text-acc (open-output-string)) ;; tool use accumulators: id -> alist (tu-table (make-hash-table)) - (current-idx (make-parameter #f))) + (current-idx (make-parameter #f)) + (usage-acc (make-hash-table))) (let ((http-status (http-post-stream url headers (json-object->string body) (lambda (event-str) @@ -490,6 +506,17 @@ ;; Block ended ((equal? event-type "content_block_stop") (current-idx #f)) + ;; Usage from message lifecycle events + ((equal? event-type "message_start") + (let ((msg (hash-get json "message"))) + (when msg + (let ((usage (hash-get msg "usage"))) + (when (and usage (hash-table? usage)) + (hash-for-each (lambda (k v) (hash-put! usage-acc k v)) usage)))))) + ((equal? event-type "message_delta") + (let ((usage (hash-get json "usage"))) + (when (and usage (hash-table? usage)) + (hash-for-each (lambda (k v) (hash-put! usage-acc k v)) usage)))) (#t (void))))))))))))) (unless (= http-status 200) (log-error logger "stream-http-error" @@ -505,10 +532,13 @@ (or (hash-get acc "name") "unknown") (or (hash-get acc "args") "{}")))) indices))) - (values content tool-calls)))) + (values content tool-calls + (list (cons 'tokens-in (or (hash-get usage-acc "input_tokens") 0)) + (cons 'tokens-out (or (hash-get usage-acc "output_tokens") 0)) + (cons 'cost 0)))))) ;; provider-stream-chat: stream text tokens to token-cb, accumulate tool calls. -;; Returns (values content-str tool-call-list). +;; Returns (values content-str tool-call-list usage-alist). (def (provider-stream-chat provider messages tools token-cb) (unless (provider-api-key provider) (error 'provider-stream-chat @@ -527,7 +557,7 @@ (content (or (message-content response) "")) (tcs (or (message-tool-calls response) '()))) (when (> (string-length content) 0) (token-cb content)) - (values content tcs))) + (values content tcs '()))) (else (error 'provider-stream-chat (format "Unknown provider '~a'" (provider-name provider)))))) --- a/src/jcode/ui/tui-input.ss +++ b/src/jcode/ui/tui-input.ss @@ -64,13 +64,20 @@ ;; Enter: newline in multi-line mode ((= key TB_KEY_ENTER) - ;; If completion is active, accept it - (if (input-state-completion inp) - (begin (accept-completion! inp) 'continue) - ;; If input is a single line with no newlines, submit - (if (string-contains (input-state-text inp) "\n") - (begin (insert-char! inp #\newline) 'continue) - 'submit))) + (cond + ;; Completion active: if text already matches selected, submit directly + ((input-state-completion inp) + (let* ((comp (input-state-completion inp)) + (selected (list-ref comp (input-state-comp-selected inp))) + (text (string-trim (input-state-text inp)))) + (if (equal? text (car selected)) + (begin (input-state-completion-set! inp #f) 'submit) + (begin (accept-completion! inp) 'continue)))) + ;; Multi-line: insert newline + ((string-contains (input-state-text inp) "\n") + (insert-char! inp #\newline) 'continue) + ;; Single line: submit + (#t 'submit))) ;; Ctrl-C: cancel ((= key TB_KEY_CTRL_C) 'cancel) @@ -237,7 +244,7 @@ (let ((comp (input-state-completion inp))) (when comp (let ((selected (list-ref comp (input-state-comp-selected inp)))) - (input-state-text-set! inp (string-append (car selected) " ")) + (input-state-text-set! inp (car selected)) (input-state-cursor-pos-set! inp (string-length (input-state-text inp))) (input-state-completion-set! inp #f))))) --- a/src/jcode/ui/tui.ss +++ b/src/jcode/ui/tui.ss @@ -548,7 +548,10 @@ (tui-stream-token! state token))) (current-tool-cb (lambda (event name args) - (tui-tool-event! state event name args)))) + (tui-tool-event! state event name args))) + (current-usage-cb + (lambda (usage) + (tui-usage-update! state usage)))) (agent-run s-id text) ;; Finalize streaming (tui-log "run-agent: agent-run returned, buf-len=~a" @@ -596,6 +599,20 @@ (loop (cdr msgs))))))) (app-state-dirty?-set! state #t)))) +(def (tui-usage-update! state usage) + "Accumulate token/cost usage from a streaming response." + (for-each + (lambda (pair) + (case (car pair) + ((tokens-in) (app-state-tokens-in-set! state + (+ (app-state-tokens-in state) (cdr pair)))) + ((tokens-out) (app-state-tokens-out-set! state + (+ (app-state-tokens-out state) (cdr pair)))) + ((cost) (app-state-cost-set! state + (+ (app-state-cost state) (cdr pair)))))) + usage) + (app-state-dirty?-set! state #t)) + (def (tool-event-metadata name args) (cond ((hash-table? args)