steeldb/agent/

needle.rs

//! Cactus **Needle** provider — drive the agent with the 26M, tool-use-native Needle model (encoder-
//! decoder, distilled for function-calling) served by the needle-code torch port. Needle is a single-turn
//! tool selector: given a `query` string + a tool list, it emits a JSON array of tool calls. We adapt
//! that to the multi-turn [`LlmProvider`] contract — the agent loop already re-invokes us with tool
//! results appended, so each call flattens the running conversation into `query`, POSTs to `/generate`,
//! and parses the returned tool-call JSON into a [`Turn`].
//!
//! Wire format (needle-code `POST /generate`):
//! ```json
//! { "query": "…task + context…", "tools": "[{\"name\":…,\"description\":…,\"parameters\":…}]",
//!   "max_gen_len": 256, "seed": 0, "constrained": true }
//! ```
//! → `{ "result": "[{\"name\":\"ikl_query\",\"arguments\":{…}}]" }`. Note tools use **Needle's own
//! schema** (`{name,description,parameters}`), NOT the OpenAI `{type,function}` envelope, and `tools` is
//! a JSON **string**. Unlike Qwen, Needle emits tool calls as structured JSON directly, so this provider
//! parses them itself and is NOT wrapped in the Qwen harness.

use crate::agent::provider::LlmProvider;
use crate::agent::types::{Block, Msg, Role, Stop, ToolSpec, Turn};
use async_trait::async_trait;
use serde_json::{json, Value};

pub struct NeedleProvider {
    client: reqwest::Client,
    base_url: String,
    max_gen_len: u32,
}

impl NeedleProvider {
    pub fn new(base_url: String) -> NeedleProvider {
        NeedleProvider { client: reqwest::Client::new(), base_url, max_gen_len: 256 }
    }
}

/// Strip the entity-linked recall appendix the agent loop injects into the opening question — useful
/// context for a chat model, but noise for Needle's compact tool router (which will echo salient phrases
/// into args).
fn core_question(t: &str) -> &str {
    match t.find("\n\n[Retrieved evidence") {
        Some(i) => t[..i].trim_end(),
        None => t.trim(),
    }
}

/// Flatten the neutral conversation into a single Needle `query` string: system framing first, then a
/// role-labelled transcript (user asks, prior assistant tool calls, and tool results), ending on a cue
/// for the next action.
pub fn build_query(system: &str, msgs: &[Msg]) -> String {
    let mut q = String::new();
    if !system.trim().is_empty() {
        q.push_str(system.trim());
        q.push_str("\n\n");
    }
    for m in msgs {
        match m.role {
            Role::User => {
                for b in &m.blocks {
                    match b {
                        Block::Text(t) => {
                            q.push_str("User: ");
                            q.push_str(core_question(t));
                            q.push('\n');
                        }
                        Block::ToolResult { content, .. } => {
                            q.push_str("Tool result: ");
                            q.push_str(content);
                            q.push('\n');
                        }
                        _ => {}
                    }
                }
            }
            Role::Assistant => {
                for b in &m.blocks {
                    match b {
                        Block::Text(t) if !t.is_empty() => {
                            q.push_str("Assistant: ");
                            q.push_str(t);
                            q.push('\n');
                        }
                        Block::ToolUse { name, input, .. } => {
                            q.push_str(&format!("Assistant called {name}({input})\n"));
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    q.trim_end().to_string()
}

/// Translate neutral tool specs into Needle's schema and serialise to the JSON **string** the server
/// expects: `[{"name","description","parameters"}]`. The schema is MINIMISED to match what the finetune
/// trained on — a terse description and enum-only properties (no prose). A finetuned 26M model keys on
/// the exact schema shape, and verbose prose both overflows context and diverges from training.
pub fn to_needle_tools(tools: &[ToolSpec]) -> String {
    let arr: Vec<Value> = tools.iter().map(|t| json!({"name": t.name, "description": minimal_desc(&t.name), "parameters": minimal_params(&t.schema)})).collect();
    Value::Array(arr).to_string()
}

fn minimal_desc(name: &str) -> String {
    match name {
        "run_workflow" => "Run one analysis over the rows.".to_string(),
        "search" => "Search the documents for a phrase.".to_string(),
        other => other.to_string(),
    }
}

/// Strip per-property `description` prose (keep `type`/`enum`/`items`) so the schema matches the
/// finetune's minimal form.
fn minimal_params(schema: &Value) -> Value {
    let mut s = schema.clone();
    if let Some(props) = s.get_mut("properties").and_then(|p| p.as_object_mut()) {
        for (_k, v) in props.iter_mut() {
            if let Some(o) = v.as_object_mut() {
                o.remove("description");
            }
        }
    }
    s
}

/// Parse Needle's `result` into a [`Turn`]. `result` is normally a JSON array of `{name, arguments}`
/// tool calls; if it isn't valid tool-call JSON, treat it as a final text answer.
pub fn parse_result(result: &str) -> Turn {
    let trimmed = result.trim();
    if let Ok(Value::Array(calls)) = serde_json::from_str::<Value>(trimmed) {
        let mut tool_uses = Vec::new();
        for (i, call) in calls.iter().enumerate() {
            let Some(name) = call.get("name").and_then(|x| x.as_str()) else { continue };
            // arguments may be an object or an escaped JSON string
            let input = match call.get("arguments") {
                Some(Value::String(s)) => serde_json::from_str(s).unwrap_or(json!({})),
                Some(v) => v.clone(),
                None => json!({}),
            };
            tool_uses.push((format!("needle-{i}"), name.to_string(), input));
        }
        if !tool_uses.is_empty() {
            return Turn { text: String::new(), tool_uses, stop: Stop::ToolUse };
        }
    }
    // Not tool calls → a plain answer.
    Turn { text: trimmed.to_string(), tool_uses: Vec::new(), stop: Stop::EndTurn }
}

#[async_trait]
impl LlmProvider for NeedleProvider {
    fn name(&self) -> &str {
        "needle"
    }

    // Needle is a tool selector, not a chat model — never trust it to phrase numbers; the agent renders
    // a deterministic template from the program result instead.
    fn synthesizes(&self) -> bool {
        false
    }

    async fn chat(&self, _system: &str, msgs: &[Msg], tools: &[ToolSpec]) -> Result<Turn, String> {
        let url = format!("{}/generate", self.base_url.trim_end_matches('/'));
        // A 26M router is reliable with a tiny, non-overlapping tool set. run_workflow already carries the
        // facet names as enums, so schema-discovery tools (list_facets/facet_tokens/ikl_query) are just
        // noise for Needle — keep only the DSL entry point and free-text search.
        const NEEDLE_TOOLS: &[&str] = &["run_workflow", "search"];
        let tools: Vec<ToolSpec> = tools.iter().filter(|t| NEEDLE_TOOLS.contains(&t.name.as_str())).cloned().collect();
        let tools = tools.as_slice();
        // Needle is a compact tool router; the verbose Qwen-oriented system prompt (with its examples and
        // MECE guidance) confuses it. Feed a lean instruction plus the conversation instead.
        // Lean framing + a couple of DSL exemplars: a 26M router benefits from seeing the exact shape of
        // a good run_workflow call. facet_* are always facet NAMES (from list_facets), never a program.
        const LEAN_SYSTEM: &str = "Call run_workflow. Your main job is `constraints`: the facet/value tokens that narrow the rows to what the question is about. Leave `program` out unless the question clearly needs a specific cross/rank/path. facet names come from the schema.\n\
Examples:\n\
Q: Tell me about electric vehicles with range over 500km.\n\
{\"name\":\"run_workflow\",\"arguments\":{\"constraints\":[\"powertrain/electric\",\"(num range_km gt 500)\"]}}\n\
Q: For each country, which powertrain is most common?\n\
{\"name\":\"run_workflow\",\"arguments\":{\"program\":\"crosstab\",\"facet_a\":\"country\",\"facet_b\":\"powertrain\"}}";
        let body = json!({
            "query": build_query(LEAN_SYSTEM, msgs),
            "tools": to_needle_tools(tools),
            "max_gen_len": self.max_gen_len,
            "seed": 0,
            "constrained": true,
        });
        let resp = self.client.post(&url).json(&body).send().await.map_err(|e| format!("needle request: {e}"))?;
        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            return Err(format!("needle {status}: {text}"));
        }
        let v: Value = resp.json().await.map_err(|e| format!("needle decode: {e}"))?;
        let result = v.get("result").and_then(|r| r.as_str()).ok_or("needle: no result field")?;
        Ok(parse_result(result))
    }
}

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

    #[test]
    fn needle_tools_use_bare_schema_not_openai_envelope() {
        let tools = vec![ToolSpec {
            name: "crosstab".into(),
            description: "cross two facets".into(),
            schema: json!({"type": "object", "properties": {"a": {"type": "string"}}}),
        }];
        let s = to_needle_tools(&tools);
        let v: Value = serde_json::from_str(&s).unwrap();
        assert_eq!(v[0]["name"], "crosstab");
        assert!(v[0].get("parameters").is_some());
        assert!(v[0].get("function").is_none(), "must NOT use the OpenAI {{type,function}} envelope");
    }

    #[test]
    fn parse_result_reads_tool_calls() {
        let t = parse_result(r#"[{"name":"crosstab","arguments":{"a":"country","b":"powertrain"}}]"#);
        assert_eq!(t.stop, Stop::ToolUse);
        assert_eq!(t.tool_uses.len(), 1);
        assert_eq!(t.tool_uses[0].1, "crosstab");
        assert_eq!(t.tool_uses[0].2["a"], "country");
    }

    #[test]
    fn parse_result_handles_escaped_arguments_string() {
        let t = parse_result(r#"[{"name":"rank","arguments":"{\"facet\":\"powertrain\"}"}]"#);
        assert_eq!(t.tool_uses[0].2["facet"], "powertrain");
    }

    #[test]
    fn parse_result_falls_back_to_text() {
        let t = parse_result("Japan is mostly electric.");
        assert_eq!(t.stop, Stop::EndTurn);
        assert!(t.tool_uses.is_empty());
        assert_eq!(t.text, "Japan is mostly electric.");
    }

    #[test]
    fn build_query_labels_turns_and_tool_results() {
        let msgs = vec![
            Msg::user_text("how many EVs per country?"),
            Msg { role: Role::Assistant, blocks: vec![Block::ToolUse { id: "1".into(), name: "crosstab".into(), input: json!({"a": "country"}) }] },
            Msg { role: Role::User, blocks: vec![Block::ToolResult { id: "1".into(), content: "japan=2".into(), is_error: false }] },
        ];
        let q = build_query("You are SteelDB.", &msgs);
        assert!(q.starts_with("You are SteelDB."));
        assert!(q.contains("User: how many EVs per country?"));
        assert!(q.contains("Assistant called crosstab"));
        assert!(q.contains("Tool result: japan=2"));
    }
}