(MOT-4358) feat(provider-groq): add the Groq provider worker

llm-router/Cargo.toml

@@ -30,7 +30,7 @@ serde_json = "1"
 # Must stay on the same schemars major as iii-sdk so the derived
 # request/response schemas match what the SDK emits at registration.
 schemars = "0.8"
-tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time", "signal"] }
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time", "signal", "fs"] }
 async-trait = "0.1"
 thiserror = "2"
 futures = "0.3"
@@ -40,6 +40,17 @@ regex = "1"
 clap = { version = "4", features = ["derive", "env"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
+# Local tokenizers for provider_scaffold::tiktoken_count (o200k/cl100k
+# vocabularies embedded in the binary — counting never touches the network).
+tiktoken-rs = "0.11"
+# A model's own vocabulary for provider_scaffold::vocabulary_count, for the
+# families that publish no metering endpoint. default-features off keeps the
+# graph pure Rust: the defaults pull onig (C) and esaxx (C++), which would
+# need a cross C toolchain on every release target.
+tokenizers = { version = "0.23", default-features = false, features = ["fancy-regex"] }
+# Fetching a vocabulary once on a cold cache. Same TLS backend as the
+# provider crates, so the graph resolves one reqwest.
+reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
 
 [dev-dependencies]
 tokio = { version = "1", features = ["test-util"] }

llm-router/README.md

@@ -58,6 +58,7 @@ partial content, so consumers never hang on a half-open stream.
 | `router::models::get` | Fetch one model record (`null` when unknown). |
 | `router::models::supports` | Check one capability flag for one model. |
 | `router::embed` | Batch text embeddings through the first embed-capable provider in the registry (`{model?, provider?, input[]}` → `{provider, model, embeddings[][]}`); providers without an embed surface are skipped. |
+| `router::count_tokens` | Count prompt tokens: `{model, provider?, system_prompt?, tools?, messages}` → `{provider, model, tokens, estimator}`, resolved with the same routing rules as `router::chat` and forwarded to `provider::<id>::count_tokens`. Never runs the model and costs nothing; `estimator` is `provider` (metering API) or `tiktoken` (local tokenizer). A provider without the surface is a typed `router/no_token_counter` error, so callers can fall back to their own estimate. |
 | `router::provider::list` | Registered providers with `configured` / `available` status. |
 | `router::system_prompt::get` | Effective identity prompt for `{provider?}` (the configured `default_provider` when omitted): operator override when set → provider-declared → `null`. `null` also when the provider is unknown — callers fall back to their own default prompt. |
 
@@ -77,8 +78,9 @@ registration token, and every later protocol call must present it.
 | `router::provider::update_credential` | Persist a refreshed credential (OAuth write-back). |
 | `router::models::reconcile` | Replace the provider's catalog slice in one write. |
 
-The provider worker itself exposes `provider::<id>::stream` and, when it
-supports model discovery, `provider::<id>::refresh_models`.
+The provider worker itself exposes `provider::<id>::stream` and, when
+capable, `provider::<id>::refresh_models` (model discovery) and
+`provider::<id>::count_tokens` (prompt token counting).
 
 ## Configuration
 
@@ -208,6 +210,13 @@ tailored to the provider's model family. The router stores it in the registry
 and serves it (unless the config slice sets an override) via
 `router::system_prompt::get`; the harness fetches it at turn creation.
 
+A provider may also register `provider::<id>::count_tokens` (`{model,
+system_prompt?, tools?, messages}` → `{model, tokens, estimator}`) to serve
+`router::count_tokens`: an exact count from a provider metering API
+(`estimator: "provider"`) or a local tokenizer estimate (`estimator:
+"tiktoken"`). Providers without it simply make `router::count_tokens` return
+a typed `router/no_token_counter` error for that provider.
+
 The first real provider implementing this protocol is
 [`provider-anthropic/`](https://github.com/iii-hq/workers/tree/main/provider-anthropic) — useful as a reference
 implementation alongside the scripted provider in the integration tests.

llm-router/src/count_tokens.rs

@@ -0,0 +1,159 @@
+//! `router::count_tokens` — one front door for prompt token counting,
+//! resolved with the SAME `decide()` the chat pipeline uses: `{model,
+//! provider?}` maps to a provider, and the request is forwarded to
+//! `provider::<id>::count_tokens`. Providers answer with an exact count from
+//! their metering API (`estimator: "provider"`) or a local tokenizer
+//! estimate (`estimator: "tiktoken"`); a provider without the surface is a
+//! typed `router/no_token_counter` error so callers can fall back to their
+//! own estimate. Counting never runs the model and costs nothing.
+
+use std::future::Future;
+use std::sync::Arc;
+
+use iii_sdk::errors::Error;
+use iii_sdk::protocol::TriggerRequest;
+use iii_sdk::IIIClient;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+
+use crate::catalog::store::CatalogStore;
+use crate::config::state::{snapshot, ConfigCell};
+use crate::registry::store::RegistryStore;
+use crate::routing::{decide, DecideInput};
+use crate::types::errors::{is_function_not_found, RouterCode, RouterError};
+use crate::types::messages::AgentMessage;
+use crate::types::model::AgentFunction;
+
+/// Counting is bounded and non-streaming, like `router::embed`.
+const COUNT_TOKENS_TIMEOUT_MS: u64 = 30_000;
+
+#[derive(Debug, Deserialize, JsonSchema)]
+pub struct RouterCountTokensRequest {
+    /// Model id the prompt targets; routes exactly like `router::chat` and
+    /// selects the provider's tokenizer. Optional only when `provider` pins
+    /// the destination.
+    #[serde(default)]
+    pub model: Option<String>,
+    /// Pin an explicit provider, bypassing heuristics (optional).
+    #[serde(default)]
+    pub provider: Option<String>,
+    /// System prompt counted as part of the request (optional).
+    #[serde(default)]
+    pub system_prompt: Option<String>,
+    /// Function invocation schemas the turn would carry; their serialized
+    /// schemas count toward the total (optional).
+    #[serde(default)]
+    pub tools: Option<Vec<AgentFunction>>,
+    /// Wire agent messages, the same shape `router::chat` accepts. Must be
+    /// non-empty.
+    pub messages: Vec<AgentMessage>,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct RouterCountTokensResponse {
+    pub provider: String,
+    pub model: String,
+    /// Prompt tokens the provider counted for the assembled request.
+    pub tokens: u64,
+    /// `provider` when a provider metering API produced the count,
+    /// `tiktoken` when a local tokenizer estimated it.
+    pub estimator: String,
+}
+
+/// What `provider::<id>::count_tokens` answers; `provider` is stamped on by
+/// the router.
+#[derive(Debug, Deserialize)]
+struct ProviderCountTokensReply {
+    model: String,
+    tokens: u64,
+    estimator: String,
+}
+
+pub fn make_count_tokens(
+    iii: IIIClient,
+    registry: Arc<RegistryStore>,
+    catalog: Arc<CatalogStore>,
+    config: ConfigCell,
+) -> impl Fn(RouterCountTokensRequest) -> BoxedCountTokensFuture + Send + Sync + 'static {
+    move |req: RouterCountTokensRequest| {
+        let (iii, registry, catalog, config) = (
+            iii.clone(),
+            registry.clone(),
+            catalog.clone(),
+            config.clone(),
+        );
+        Box::pin(async move {
+            // Providers stay dumb pipes: the router never injects placeholder
+            // messages to make an empty request countable.
+            if req.messages.is_empty() {
+                return Err(RouterError::new(
+                    RouterCode::InvalidRequest,
+                    "messages must not be empty",
+                )
+                .into());
+            }
+            let model = req.model.unwrap_or_default();
+            if model.is_empty() && req.provider.is_none() {
+                return Err(RouterError::new(
+                    RouterCode::InvalidRequest,
+                    "model or provider is required",
+                )
+                .into());
+            }
+            // Same inputs, same decide(), same error codes as the chat
+            // pipeline's routing step.
+            let config = snapshot(&config);
+            let heuristics = config.settings().routing_heuristics.clone();
+            let default_provider = config.settings().default_provider.clone();
+            let candidates = decide(&DecideInput {
+                model: model.clone(),
+                provider: req.provider,
+                registered_providers: registry.ids().await,
+                catalog: catalog.model_ids().await,
+                heuristics,
+                default_provider,
+            })
+            .map_err(Error::from)?;
+            let provider = candidates[0].clone();
+
+            let reply = iii
+                .trigger(TriggerRequest {
+                    function_id: format!("provider::{provider}::count_tokens"),
+                    payload: json!({
+                        "model": model,
+                        "system_prompt": req.system_prompt,
+                        "tools": req.tools,
+                        "messages": req.messages,
+                    }),
+                    action: None,
+                    timeout_ms: Some(COUNT_TOKENS_TIMEOUT_MS),
+                })
+                .await
+                .map_err(|e| {
+                    if is_function_not_found(&e) {
+                        Error::Handler(format!(
+                            "router/no_token_counter: provider '{provider}' has no token counter"
+                        ))
+                    } else {
+                        e
+                    }
+                })?;
+            let reply: ProviderCountTokensReply = serde_json::from_value(reply).map_err(|e| {
+                Error::Handler(format!(
+                    "router/bad_provider_response: provider::{provider}::count_tokens \
+                         returned an invalid response: {e}"
+                ))
+            })?;
+            Ok(RouterCountTokensResponse {
+                provider,
+                model: reply.model,
+                tokens: reply.tokens,
+                estimator: reply.estimator,
+            })
+        })
+    }
+}
+
+type BoxedCountTokensFuture =
+    std::pin::Pin<Box<dyn Future<Output = Result<RouterCountTokensResponse, Error>> + Send>>;

llm-router/src/embed.rs

@@ -17,6 +17,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
 
 use crate::registry::store::RegistryStore;
+use crate::types::errors::is_function_not_found;
 
 const EMBED_TIMEOUT_MS: u64 = 30_000;
 
@@ -59,7 +60,7 @@ async fn try_provider(
         .await;
     let reply = match reply {
         Ok(v) => v,
-        Err(e) if e.to_string().contains("function_not_found") => return Ok(None),
+        Err(e) if is_function_not_found(&e) => return Ok(None),
         Err(e) => return Err(e),
     };
     let model = reply

llm-router/src/lib.rs

@@ -5,6 +5,7 @@ pub mod catalog;
 pub mod channels;
 pub mod chat;
 pub mod config;
+pub mod count_tokens;
 pub mod embed;
 pub mod manifest;
 pub mod provider_scaffold;

llm-router/src/provider_scaffold/chat_framing.rs

@@ -0,0 +1,134 @@
+//! How an assembled chat request is walked for counting, independent of which
+//! tokenizer does the counting.
+//!
+//! Every local counter — tiktoken for the OpenAI families, a model's own
+//! vocabulary for everything else — has to agree on the same three things:
+//! which text a message contributes to the wire, what the chat framing costs
+//! on top of that text, and how a tool schema is serialized. Only the
+//! encoding differs. Keeping the walk here means a new tokenizer is one
+//! closure rather than a second copy of these rules, and a fix to the walk
+//! reaches every counter at once.
+
+use serde_json::json;
+
+use crate::provider_scaffold::names::encode_tool_name;
+use crate::types::content::ContentBlock;
+use crate::types::messages::AgentMessage;
+use crate::types::model::AgentFunction;
+
+/// Chat-framing overhead per wire message row (role tag + separators),
+/// OpenAI's published ~4-tokens-per-message heuristic for ChatML-framed
+/// conversations. The system prompt is one such row.
+pub const TOKENS_PER_MESSAGE: u64 = 4;
+
+/// Reply priming the API appends to every prompt (the assistant start tag).
+pub const TOKENS_REPLY_PRIMING: u64 = 2;
+
+/// The text a message contributes to the wire: text blocks, plus each
+/// function call's encoded name and serialized arguments, plus function
+/// result bodies. Thinking blocks are dropped (never replayed on this wire),
+/// images are skipped (their token cost is model-specific, not textual), and
+/// custom messages never reach the provider.
+pub fn message_text(message: &AgentMessage) -> Option<String> {
+    let mut parts: Vec<String> = Vec::new();
+    let content = match message {
+        AgentMessage::User(m) => &m.content,
+        AgentMessage::Assistant(m) => &m.content,
+        AgentMessage::FunctionResult(m) => &m.content,
+        AgentMessage::Custom(_) => return None,
+    };
+    for block in content {
+        match block {
+            ContentBlock::Text { text } => parts.push(text.clone()),
+            ContentBlock::FunctionCall {
+                function_id,
+                arguments,
+                ..
+            } => {
+                parts.push(encode_tool_name(function_id));
+                parts.push(arguments.to_string());
+            }
+            _ => {}
+        }
+    }
+    Some(parts.join("\n"))
+}
+
+/// The serialized form of one tool schema, as it reaches the wire.
+pub fn tool_schema_text(tool: &AgentFunction) -> String {
+    json!({
+        "name": encode_tool_name(&tool.name),
+        "description": tool.description,
+        "parameters": tool.parameters,
+    })
+    .to_string()
+}
+
+/// An assembled request reduced to the text that gets counted: one entry per
+/// framed wire row (system prompt, each message), and one per tool schema.
+/// Schemas are listed apart because they ride inside the request rather than
+/// as rows of their own, so they carry no per-message framing.
+pub struct Framed {
+    pub rows: Vec<String>,
+    pub schemas: Vec<String>,
+}
+
+/// Reduce a request to what a counter has to encode. Splitting this from the
+/// arithmetic is what lets a remote tokenizer — one that takes text rather
+/// than messages — batch the whole request into a single call.
+pub fn frame(
+    system_prompt: Option<&str>,
+    tools: &[AgentFunction],
+    messages: &[AgentMessage],
+) -> Framed {
+    let mut rows = Vec::new();
+    if let Some(system) = system_prompt.filter(|s| !s.is_empty()) {
+        rows.push(system.to_string());
+    }
+    rows.extend(messages.iter().filter_map(message_text));
+    Framed {
+        rows,
+        schemas: tools.iter().map(tool_schema_text).collect(),
+    }
+}
+
+impl Framed {
+    /// Every counted text in one string, for a tokenizer that is reached over
+    /// the network and should be called once rather than per row. Costs the
+    /// separator tokens at each join, which is a handful of tokens against a
+    /// whole request.
+    pub fn joined(&self) -> String {
+        self.rows
+            .iter()
+            .chain(self.schemas.iter())
+            .cloned()
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+
+    /// The request total once the counted text has come to `text_tokens`:
+    /// reply priming plus this request's framing on top.
+    pub fn total_from(&self, text_tokens: u64) -> u64 {
+        TOKENS_REPLY_PRIMING + self.rows.len() as u64 * TOKENS_PER_MESSAGE + text_tokens
+    }
+}
+
+/// Count an assembled chat request with `count_text` as the encoder: reply
+/// priming, plus one framed row per system prompt and message, plus each
+/// tool's serialized schema. The framing constants are the caller-independent
+/// part; `count_text` is the only thing a tokenizer changes.
+pub fn count_framed_chat(
+    system_prompt: Option<&str>,
+    tools: &[AgentFunction],
+    messages: &[AgentMessage],
+    count_text: impl Fn(&str) -> u64,
+) -> u64 {
+    let framed = frame(system_prompt, tools, messages);
+    let text_tokens: u64 = framed
+        .rows
+        .iter()
+        .chain(framed.schemas.iter())
+        .map(|text| count_text(text))
+        .sum();
+    framed.total_from(text_tokens)
+}