Adding usage data printout

include/json_types.hpp

@@ -81,10 +81,18 @@ inline void to_json(nlohmann::json& j, const Choice& c) {
     if (!c.finish_reason.empty()) j["finish_reason"] = c.finish_reason;
 }
 
+struct Usage {
+    int completion_tokens;
+    int prompt_tokens;
+    int total_tokens;
+};
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Usage, completion_tokens, prompt_tokens, total_tokens)
+
 struct ChatResponse {
     std::vector<Choice> choices;
+    Usage usage;
 };
-NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(ChatResponse, choices)
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(ChatResponse, choices, usage)
 
 struct ModelInfo {
     std::string name;

src/ollama_client.cpp

@@ -94,8 +94,12 @@ ChatResponse OllamaClient::chat(const std::vector<Message>& messages,
     std::string response_buf = http_request(url, "POST", body_str);
 
     try {
+        auto ret = nlohmann::json::parse(response_buf).get<ChatResponse>();
         std::cerr << "Repsonse: " << response_buf << std::endl;
-        return nlohmann::json::parse(response_buf).get<ChatResponse>();
+        std::cerr << "Total Tokens: " << ret.usage.total_tokens << std::endl;
+        std::cerr << "Input Tokens: " << ret.usage.prompt_tokens << std::endl;
+        std::cerr << "Generated Tokens: " << ret.usage.completion_tokens << std::endl;
+        return ret;
     } catch (const nlohmann::json::exception& e) {
         throw OllamaClientError(std::string("JSON parse error: ") + e.what());
     }

tests/run_test.sh

@@ -25,24 +25,24 @@ function tool_call() {
     name="$1"
     args="$2"
     args_string=`sanitize_json "$2"`
-    echo '{"choices":[{"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"'$name'","arguments":"'"$args_string"'"}}],"content":null},"finish_reason":"tool_calls"}]}'
+    echo '{"choices":[{"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"'$name'","arguments":"'"$args_string"'"}}],"content":null},"finish_reason":"tool_calls"}],"usage":{"completion_tokens":336,"prompt_tokens":14,"total_tokens":350}}'
 }
 function be_done() {
     content=`sanitize_json "$1"`
 
-    echo '{"choices":[{"message":{"role":"assistant","content":"'"$content"'"},"finish_reason":"stop"}]}'
+    echo '{"choices":[{"message":{"role":"assistant","content":"'"$content"'"},"finish_reason":"stop"}],"usage":{"completion_tokens":336,"prompt_tokens":14,"total_tokens":350}}'
 }
 
 function reasoning_step() {
     content=`sanitize_json "$1"`
 
-    echo '{"choices":[{"message":{"role":"assistant","content":"","reasoning":"'"$content"'"},"finish_reason":"stop"}]}'
+    echo '{"choices":[{"message":{"role":"assistant","content":"","reasoning":"'"$content"'"},"finish_reason":"stop"}],"usage":{"completion_tokens":336,"prompt_tokens":14,"total_tokens":350}}'
 }
 
 function reasoning_content_step() {
     content=`sanitize_json "$1"`
 
-    echo '{"choices":[{"message":{"role":"assistant","content":"","reasoning_content":"'"$content"'"},"finish_reason":"stop"}]}'
+    echo '{"choices":[{"message":{"role":"assistant","content":"","reasoning_content":"'"$content"'"},"finish_reason":"stop"}],"usage":{"completion_tokens":336,"prompt_tokens":14,"total_tokens":350}}'
 }
 
 export -f check_tool_call

tests/test_smoke_basic.sh

@@ -1,9 +1,7 @@
 #!/bin/bash
 # test_smoke_basic.sh - Verifies drharness outputs a mock response correctly
 
-MOCK_RESPONSE='{"choices":[{"message":{"role":"assistant","content":"Test response from mock"},"finish_reason":"stop"}]}'
-
-OUTPUT=$(echo "${MOCK_RESPONSE}" | "${DRHARNESS}" -F chat.json -n -m "Hi there" -d)
+OUTPUT=$(be_done 'Test response from mock' | "${DRHARNESS}" -F chat.json -n -m "Hi there" -d)
 
 if [ "${OUTPUT}" != "Test response from mock" ]; then
     echo "FAIL: smoke_basic - expected 'Test response from mock', got '${OUTPUT}'"