From 103911ad16323805f2308046d99b7896fea05112 Mon Sep 17 00:00:00 2001 From: Dale Seo Date: Thu, 10 Jul 2025 10:36:32 -0400 Subject: [PATCH 1/2] deps: add tiktoken-rs --- Cargo.lock | 43 +++++++++++++++++++++++++++++ crates/apollo-mcp-server/Cargo.toml | 1 + 2 files changed, 44 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 86b24e77..ce5474aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -216,6 +216,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.12", + "tiktoken-rs", "tokio", "tokio-util", "tracing", @@ -367,6 +368,21 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + [[package]] name = "bitflags" version = "1.3.2" @@ -420,6 +436,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "234113d19d0d7d613b40e86fb654acf958910802bcceab913a4f9e7cda03b1a4" dependencies = [ "memchr", + "regex-automata 0.4.9", "serde", ] @@ -777,6 +794,17 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata 0.4.9", + "regex-syntax 0.8.5", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -2659,6 +2687,21 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tiktoken-rs" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625" +dependencies = [ + "anyhow", + "base64", + "bstr", + "fancy-regex", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.41" diff --git a/crates/apollo-mcp-server/Cargo.toml b/crates/apollo-mcp-server/Cargo.toml index 3997ca7f..62663ba1 100644 --- a/crates/apollo-mcp-server/Cargo.toml +++ b/crates/apollo-mcp-server/Cargo.toml @@ -26,6 +26,7 @@ rmcp = { version = "0.2", features = [ serde.workspace = true serde_json.workspace = true thiserror.workspace = true +tiktoken-rs = "0.7.0" tokio.workspace = true tracing.workspace = true tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } From 9cd3497b2826940f4be192b506b9835e2e1ff1f2 Mon Sep 17 00:00:00 2001 From: Dale Seo Date: Thu, 10 Jul 2025 12:31:59 -0400 Subject: [PATCH 2/2] feat: better token counting --- crates/apollo-mcp-server/src/lib.rs | 1 + crates/apollo-mcp-server/src/operations.rs | 14 +-- .../apollo-mcp-server/src/token_counting.rs | 93 +++++++++++++++++++ 3 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 crates/apollo-mcp-server/src/token_counting.rs diff --git a/crates/apollo-mcp-server/src/lib.rs b/crates/apollo-mcp-server/src/lib.rs index 93a3c029..655edef5 100644 --- a/crates/apollo-mcp-server/src/lib.rs +++ b/crates/apollo-mcp-server/src/lib.rs @@ -9,3 +9,4 @@ pub mod operations; pub mod sanitize; pub(crate) mod schema_tree_shake; pub mod server; +pub(crate) mod token_counting; diff --git a/crates/apollo-mcp-server/src/operations.rs b/crates/apollo-mcp-server/src/operations.rs index d67d6eb5..8bde252b 100644 --- a/crates/apollo-mcp-server/src/operations.rs +++ b/crates/apollo-mcp-server/src/operations.rs @@ -3,6 +3,7 @@ use crate::errors::{McpError, OperationError}; use crate::event::Event; use crate::graphql::{self, OperationDetails}; use crate::schema_tree_shake::{DepthLimit, SchemaTreeShaker}; +use crate::token_counting; use apollo_compiler::ast::{Document, OperationType, Selection}; use apollo_compiler::schema::ExtendedType; use apollo_compiler::validation::Valid; @@ -566,12 +567,13 @@ impl Operation { ); let character_count = tool_character_length(&tool); match character_count { - Ok(length) => info!( - "Tool {} loaded with a character count of {}. Estimated tokens: {}", - operation_name, - length, - length / 4 // We don't know the tokenization algorithm, so we just use 4 characters per token as a rough estimate. https://docs.anthropic.com/en/docs/resources/glossary#tokens - ), + Ok(length) => { + let token_estimates = token_counting::count_tokens_from_tool(&tool); + info!( + "Tool {} loaded with a character count of {}. Estimated tokens: {}", + operation_name, length, token_estimates + ); + } Err(_) => info!( "Tool {} loaded with an unknown character count", operation_name diff --git a/crates/apollo-mcp-server/src/token_counting.rs b/crates/apollo-mcp-server/src/token_counting.rs new file mode 100644 index 00000000..0293037b --- /dev/null +++ b/crates/apollo-mcp-server/src/token_counting.rs @@ -0,0 +1,93 @@ +use rmcp::model::Tool; +use rmcp::serde_json; +use std::fmt; +use tiktoken_rs::{cl100k_base, o200k_base, p50k_base}; + +#[derive(Debug, Clone)] +pub struct TokenEstimates { + pub anthropic: Option, + pub gemini: Option, + pub openai: Option, + pub fallback: usize, +} + +impl fmt::Display for TokenEstimates { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut estimates = Vec::new(); + + if let Some(count) = self.anthropic { + estimates.push(format!("{count} Anthropic tokens")); + } + if let Some(count) = self.gemini { + estimates.push(format!("{count} Gemini tokens")); + } + if let Some(count) = self.openai { + estimates.push(format!("{count} OpenAI tokens")); + } + + if estimates.is_empty() { + write!(f, "~{} tokens (fallback estimate)", self.fallback) + } else { + write!(f, "{}", estimates.join(", ")) + } + } +} + +pub fn count_tokens_from_tool(tool: &Tool) -> TokenEstimates { + let tokenizer = TokenCounter; + let tool_text = format!( + "{}\n{}\n{}", + tool.name, + tool.description.as_ref().map(|d| d.as_ref()).unwrap_or(""), + serde_json::to_string_pretty(&tool.input_schema).unwrap_or_default() + ); + tokenizer.count_tokens(&tool_text) +} + +struct TokenCounter; + +impl TokenCounter { + pub fn count_tokens(&self, text: &str) -> TokenEstimates { + let fallback = self.estimate_tokens(text); + TokenEstimates { + anthropic: self.count_anthropic_tokens(text), + gemini: self.count_gemini_tokens(text), + openai: self.count_openai_tokens(text), + fallback, + } + } + + fn count_openai_tokens(&self, text: &str) -> Option { + // Start with o200k_base (GPT-4o, o1 models) + if let Ok(tokenizer) = o200k_base() { + return Some(tokenizer.encode_with_special_tokens(text).len()); + } + + // Fallback to cl100k_base (ChatGPT, GPT-4) + if let Ok(tokenizer) = cl100k_base() { + return Some(tokenizer.encode_with_special_tokens(text).len()); + } + + // Final fallback to p50k_base (GPT-3.5, Codex) + if let Ok(tokenizer) = p50k_base() { + return Some(tokenizer.encode_with_special_tokens(text).len()); + } + + None + } + + // TODO: Implement using Anthropic's SDK or REST API (https://docs.anthropic.com/en/docs/build-with-claude/token-counting) + fn count_anthropic_tokens(&self, _text: &str) -> Option { + None + } + + // TODO: Implement their Gemini's SDK or REST API (https://ai.google.dev/api/tokens#v1beta.models.countTokens) + fn count_gemini_tokens(&self, _text: &str) -> Option { + None + } + + fn estimate_tokens(&self, text: &str) -> usize { + let character_count = text.chars().count(); + character_count / 4 + } +}