cortiq-gateway 0.2.48

//! Tool-call hygiene for OpenAI-compatible responses.
//!
//! Qwen-family chat templates express a function call as literal markup in the
//! generated text:
//!
//! ```text
//! <tool_call>
//! {"name": "get_weather", "arguments": {"city": "Paris"}}
//! </tool_call>
//! ```
//!
//! Runtimes handle that inconsistently. Some parse it into structured
//! `tool_calls` **and** still stream the raw markup as `content` — the client
//! then sees the call twice, once as prose it must not show and once as the real
//! thing, and tool-loop implementations (Open WebUI's among them) get confused.
//! Others never parse it at all, so a perfectly good call reaches the client as
//! text no tool loop will ever execute.
//!
//! This module makes the gateway the place where that is settled: markup is
//! stripped from the text, and when nothing structured came with it the markup
//! is promoted into a real `tool_calls` array.

const OPEN: &str = "<tool_call>";
const CLOSE: &str = "</tool_call>";

/// Longest suffix of `s` that is a proper prefix of `pat`. Used to hold back the
/// tail of a streamed chunk that might be the beginning of a marker split across
/// chunk boundaries (`"<tool"` + `"_call>"`).
fn partial_marker_len(s: &str, pat: &str) -> usize {
    let max = pat.len().saturating_sub(1).min(s.len());
    (1..=max)
        .rev()
        .find(|&n| {
            s.is_char_boundary(s.len() - n)
                && pat.as_bytes().starts_with(&s.as_bytes()[s.len() - n..])
        })
        .unwrap_or(0)
}

/// Incremental `<tool_call>` filter for SSE deltas.
///
/// Feed it content fragments in order; it returns the text that is safe to
/// forward and keeps the markup payloads for [`Self::captured`].
#[derive(Default)]
pub struct ToolMarkupFilter {
    /// Text held back: either an unterminated marker prefix or the inside of an
    /// open `<tool_call>` block.
    pending: String,
    inside: bool,
    captured: Vec<String>,
}

impl ToolMarkupFilter {
    pub fn new() -> Self {
        Self::default()
    }

    /// Consume one content fragment; returns the part that may be forwarded now.
    pub fn push(&mut self, chunk: &str) -> String {
        self.pending.push_str(chunk);
        let mut out = String::new();
        loop {
            if self.inside {
                let Some(end) = self.pending.find(CLOSE) else {
                    // still inside the block — nothing to forward
                    return out;
                };
                let payload: String = self.pending[..end].to_string();
                self.captured.push(payload);
                self.pending = self.pending[end + CLOSE.len()..].to_string();
                self.inside = false;
                continue;
            }
            match self.pending.find(OPEN) {
                Some(start) => {
                    out.push_str(&self.pending[..start]);
                    self.pending = self.pending[start + OPEN.len()..].to_string();
                    self.inside = true;
                }
                None => {
                    // Hold back a trailing partial marker so a split like
                    // "…<tool" / "_call>…" is not forwarded as visible text.
                    let hold = partial_marker_len(&self.pending, OPEN);
                    let cut = self.pending.len() - hold;
                    out.push_str(&self.pending[..cut]);
                    self.pending = self.pending[cut..].to_string();
                    return out;
                }
            }
        }
    }

    /// End of stream: an unterminated block was never a tool call after all, so
    /// give the text back rather than swallowing output.
    pub fn finish(&mut self) -> String {
        let mut out = String::new();
        if self.inside {
            out.push_str(OPEN);
            self.inside = false;
        }
        out.push_str(&std::mem::take(&mut self.pending));
        out
    }

    /// Payloads found between `<tool_call>` and `</tool_call>`.
    pub fn captured(&self) -> &[String] {
        &self.captured
    }
}

/// Strip every complete `<tool_call>` block from a finished string.
/// Returns the cleaned text and the payloads that were removed.
pub fn split_tool_markup(text: &str) -> (String, Vec<String>) {
    let mut f = ToolMarkupFilter::new();
    let mut out = f.push(text);
    out.push_str(&f.finish());
    (out, f.captured)
}

/// Turn `<tool_call>` payloads into OpenAI `tool_calls` entries, numbered from
/// `start_index`. Payloads that are not a JSON object with a `name` are skipped —
/// a stray marker must not become a bogus call.
pub fn markup_to_tool_calls(payloads: &[String], start_index: usize) -> Vec<serde_json::Value> {
    let mut out = Vec::new();
    for payload in payloads {
        let Ok(v) = serde_json::from_str::<serde_json::Value>(payload.trim()) else {
            continue;
        };
        let Some(name) = v.get("name").and_then(|n| n.as_str()) else {
            continue;
        };
        // `arguments` is a JSON string on the wire; templates emit an object.
        let arguments = match v.get("arguments").or_else(|| v.get("parameters")) {
            Some(serde_json::Value::String(s)) => s.clone(),
            Some(other) => other.to_string(),
            None => "{}".to_string(),
        };
        out.push(serde_json::json!({
            "index": start_index + out.len(),
            "id": new_call_id(),
            "type": "function",
            "function": { "name": name, "arguments": arguments },
        }));
    }
    out
}

pub fn new_call_id() -> String {
    format!("call_{}", crate::admin::random_token(12))
}

/// Bring one streamed `tool_calls` array up to the OpenAI shape.
///
/// `seen` tracks which choice-local indices already had their opening delta, so
/// the identity fields (`id`, `type`) are added once and continuation deltas —
/// which legitimately carry nothing but an argument fragment — are left alone.
pub fn normalize_delta_tool_calls(
    raw: &[serde_json::Value],
    seen: &mut std::collections::HashSet<u64>,
) -> Vec<serde_json::Value> {
    let mut out = Vec::new();
    for (i, tc) in raw.iter().enumerate() {
        let index = tc.get("index").and_then(|v| v.as_u64()).unwrap_or(i as u64);
        let mut entry = serde_json::Map::new();
        entry.insert("index".into(), serde_json::json!(index));

        let first = seen.insert(index);
        if let Some(id) = tc
            .get("id")
            .and_then(|v| v.as_str())
            .filter(|s| !s.is_empty())
        {
            entry.insert("id".into(), serde_json::json!(id));
        } else if first {
            entry.insert("id".into(), serde_json::json!(new_call_id()));
        }
        if first {
            let kind = tc
                .get("type")
                .and_then(|v| v.as_str())
                .unwrap_or("function");
            entry.insert("type".into(), serde_json::json!(kind));
        }

        if let Some(f) = tc.get("function") {
            let mut func = serde_json::Map::new();
            if let Some(name) = f.get("name").and_then(|v| v.as_str()) {
                func.insert("name".into(), serde_json::json!(name));
            } else if first {
                func.insert("name".into(), serde_json::json!(""));
            }
            match f.get("arguments") {
                // Arguments are a string on the wire; some runtimes send the object.
                Some(serde_json::Value::String(s)) => {
                    func.insert("arguments".into(), serde_json::json!(s));
                }
                Some(serde_json::Value::Null) | None => {
                    if first {
                        func.insert("arguments".into(), serde_json::json!(""));
                    }
                }
                Some(other) => {
                    func.insert("arguments".into(), serde_json::json!(other.to_string()));
                }
            }
            entry.insert("function".into(), serde_json::Value::Object(func));
        }
        out.push(serde_json::Value::Object(entry));
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    const CALL: &str = r#"<tool_call>
{"name": "get_weather", "arguments": {"city": "Paris"}}
</tool_call>"#;

    #[test]
    fn markup_is_stripped_from_finished_text() {
        let (clean, payloads) = split_tool_markup(&format!("Sure!{CALL}"));
        assert_eq!(clean, "Sure!");
        assert_eq!(payloads.len(), 1);
    }

    #[test]
    fn markup_split_across_stream_chunks_never_leaks() {
        let mut f = ToolMarkupFilter::new();
        let mut seen = String::new();
        // the marker itself is broken across three fragments
        for chunk in [
            "Let me check. <tool",
            "_call>",
            "{\"name\":\"get_weather\",",
            "\"arguments\":{\"city\":\"Paris\"}}",
            "</tool",
            "_call>",
            " done",
        ] {
            seen.push_str(&f.push(chunk));
        }
        seen.push_str(&f.finish());
        assert_eq!(seen, "Let me check.  done");
        assert!(!seen.contains("tool_call"));
        assert_eq!(f.captured().len(), 1);
    }

    #[test]
    fn captured_markup_becomes_a_real_tool_call() {
        let (_, payloads) = split_tool_markup(CALL);
        let calls = markup_to_tool_calls(&payloads, 0);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0]["type"], "function");
        assert_eq!(calls[0]["index"], 0);
        assert_eq!(calls[0]["function"]["name"], "get_weather");
        // arguments must be a STRING, not an object
        assert_eq!(
            calls[0]["function"]["arguments"].as_str().unwrap(),
            r#"{"city":"Paris"}"#
        );
        assert!(calls[0]["id"].as_str().unwrap().starts_with("call_"));
    }

    #[test]
    fn a_stray_marker_does_not_become_a_call() {
        let calls = markup_to_tool_calls(&["not json at all".into(), "{\"no_name\":1}".into()], 0);
        assert!(calls.is_empty());
    }

    #[test]
    fn an_unterminated_block_is_given_back_as_text() {
        let mut f = ToolMarkupFilter::new();
        let mut out = f.push("thinking <tool_call>{\"name\":\"x\"");
        out.push_str(&f.finish());
        assert_eq!(out, "thinking <tool_call>{\"name\":\"x\"");
        assert!(f.captured().is_empty());
    }

    #[test]
    fn delta_tool_calls_get_identity_once_and_object_arguments_become_strings() {
        let mut seen = std::collections::HashSet::new();
        let opening = normalize_delta_tool_calls(
            &[serde_json::json!({"index": 0, "function": {"name": "f", "arguments": {"a": 1}}})],
            &mut seen,
        );
        assert_eq!(opening[0]["type"], "function");
        assert!(opening[0]["id"].as_str().unwrap().starts_with("call_"));
        assert_eq!(opening[0]["function"]["arguments"], r#"{"a":1}"#);

        // a continuation delta must not be handed a second id
        let cont = normalize_delta_tool_calls(
            &[serde_json::json!({"index": 0, "function": {"arguments": "\"}"}})],
            &mut seen,
        );
        assert!(cont[0].get("id").is_none());
        assert!(cont[0].get("type").is_none());
    }
}