Include token counts for Piebald that were accidentally ignored

[mike1858]

Summary by CodeRabbit

• New Features
  • Implemented tool call count tracking in conversation analytics to provide deeper insights into tool usage within conversations.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes add tool call count tracking to the Piebald analyzer by introducing a new database query function and integrating per-message tool call counts into the message conversion and statistics pipeline.

Changes

Cohort / File(s)	Summary
Tool Call Count Tracking src/analyzers/piebald.rs	Added query_tool_call_counts function to retrieve per-message tool call counts from the database via a join with message_part_tool_call. Updated convert_messages signature to accept a tool_call_counts map parameter and populate tool_calls in the Stats struct. Integrated tool call count queries into the PiebaldAnalyzer flow and parse_source to pass counts through the conversion pipeline.

Sequence Diagram

sequenceDiagram
    participant Analyzer as PiebaldAnalyzer
    participant Query as query_tool_call_counts()
    participant DB as Database
    participant Convert as convert_messages()
    participant Stats as Stats Struct

    Analyzer->>Query: Request tool call counts
    Query->>DB: JOIN message_parts with message_part_tool_call
    DB-->>Query: Return HashMap<message_id, count>
    Query-->>Analyzer: tool_call_counts map
    
    Analyzer->>Convert: Pass messages + tool_call_counts
    Convert->>Convert: Lookup count for each message
    Convert->>Stats: Populate tool_calls field (default 0 if absent)
    Stats-->>Convert: Updated Stats with tool call data
    Convert-->>Analyzer: ConversationMessage with tool_calls

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 Hark! A new counting spell takes flight,
Tool calls tallied, tracked just right,
Through databases deep, our queries do hop,
With stats made complete, the tracking won't stop! 🔢✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title claims to include 'token counts' but the summary shows the PR actually adds 'tool_calls' data, not token counts. This is misleading about what the changeset contains.	Update the title to accurately reflect that the PR adds tool call counts to Piebald analyzer, not token counts.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/add-tool-call-count-to-piebald-analyzer

📝 Coding Plan

• Generate coding plan for human review comments

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

src/analyzers/piebald.rs (1)

310-366: Consider adding test coverage for tool call count functionality.

The existing tests don't cover query_tool_call_counts or the tool call integration in convert_messages. While the existing query functions (query_chats, query_messages) also lack direct unit tests, adding an in-memory SQLite test for the new functionality would improve confidence:

💡 Optional: Test skeleton for tool call counts

#[test]
fn test_convert_messages_with_tool_calls() {
    let chats = vec![PiebaldChat {
        id: 1,
        title: Some("Test".to_string()),
        model: Some("claude-3".to_string()),
        current_directory: Some("/test".to_string()),
    }];
    
    let messages = vec![PiebaldMessage {
        id: 100,
        parent_chat_id: 1,
        role: "assistant".to_string(),
        model: Some("claude-3".to_string()),
        input_tokens: Some(50),
        output_tokens: Some(100),
        reasoning_tokens: None,
        cache_read_tokens: None,
        cache_write_tokens: None,
        created_at: "2025-12-10T14:30:00Z".to_string(),
        updated_at: "2025-12-10T14:30:00Z".to_string(),
    }];
    
    let mut tool_call_counts = HashMap::new();
    tool_call_counts.insert(100i64, 3u32);
    
    let result = convert_messages(&chats, messages, &tool_call_counts);
    
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].stats.tool_calls, 3);
}

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/analyzers/piebald.rs` around lines 310 - 366, Add unit test coverage for
tool-call counting by creating an in-memory test that constructs PiebaldChat and
PiebaldMessage instances, builds a HashMap<i64,u32> mapping message IDs to
counts (the shape returned by query_tool_call_counts), calls
convert_messages(&chats, messages, &tool_call_counts) and asserts the resulting
MessageStats.tool_calls equals the provided count; reference PiebaldChat,
PiebaldMessage, convert_messages and query_tool_call_counts when locating code
to test and place the new test alongside the existing #[cfg(test)] mod tests.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@src/analyzers/piebald.rs`:
- Around line 310-366: Add unit test coverage for tool-call counting by creating
an in-memory test that constructs PiebaldChat and PiebaldMessage instances,
builds a HashMap<i64,u32> mapping message IDs to counts (the shape returned by
query_tool_call_counts), calls convert_messages(&chats, messages,
&tool_call_counts) and asserts the resulting MessageStats.tool_calls equals the
provided count; reference PiebaldChat, PiebaldMessage, convert_messages and
query_tool_call_counts when locating code to test and place the new test
alongside the existing #[cfg(test)] mod tests.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 7c59fbc9-6eea-4313-a35d-c22e5aa0641d

📥 Commits

Reviewing files that changed from the base of the PR and between 1552a57 and a551848.

📒 Files selected for processing (1)

• src/analyzers/piebald.rs