geographdb-core 0.5.4

//! Tool schema parsing and subgraph construction for structure decoding.
//!
//! Extracts JSON function schemas from tool-use datasets (e.g. Glaive) and
//! builds small anchor/argument subgraphs inside the geometric graph.

use crate::algorithms::four_d::{GraphNode4D, GraphProperties, TemporalEdge};
use anyhow::{Context, Result};
use glam::Vec3;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap, HashSet};

/// A parsed function-calling tool schema.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolSchema {
    pub name: String,
    #[serde(default)]
    pub description: String,
    #[serde(default)]
    pub parameters: serde_json::Value,
}

impl ToolSchema {
    /// Required parameter names according to the schema.
    pub fn required_params(&self) -> Vec<String> {
        if let Some(arr) = self.parameters.get("required").and_then(|v| v.as_array()) {
            arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Parameter name → type/description map.
    pub fn param_properties(&self) -> BTreeMap<String, serde_json::Value> {
        if let Some(obj) = self
            .parameters
            .get("properties")
            .and_then(|v| v.as_object())
        {
            obj.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
        } else {
            BTreeMap::new()
        }
    }
}

/// Extract the first top-level JSON object from a text block.
///
/// Glaive tool schemas are embedded inside a `SYSTEM:` preamble; this function
/// finds the first `{`, tracks brace balance, and returns the matching object
/// string.
pub fn extract_json_object(text: &str) -> Option<String> {
    let start = text.find('{')?;
    let mut depth = 0i32;
    let mut in_string = false;
    let mut escape = false;
    let bytes = text.as_bytes();

    for i in start..bytes.len() {
        let b = bytes[i];
        if in_string {
            if escape {
                escape = false;
                continue;
            }
            if b == b'\\' {
                escape = true;
                continue;
            }
            if b == b'"' {
                in_string = false;
            }
            continue;
        }

        match b {
            b'"' => in_string = true,
            b'{' => depth += 1,
            b'}' => {
                depth -= 1;
                if depth == 0 {
                    return String::from_utf8(bytes[start..=i].to_vec()).ok();
                }
            }
            _ => {}
        }
    }
    None
}

/// Parse tool schemas from a list of `system` text snippets.
pub fn parse_tool_schemas(system_texts: &[String]) -> Vec<ToolSchema> {
    let mut seen = HashSet::new();
    let mut out = Vec::new();
    for text in system_texts {
        let Some(json) = extract_json_object(text) else {
            continue;
        };
        let Ok(schema): Result<ToolSchema, _> = serde_json::from_str(&json) else {
            continue;
        };
        if schema.name.is_empty() || seen.contains(&schema.name) {
            continue;
        }
        seen.insert(schema.name.clone());
        out.push(schema);
    }
    out
}

/// Unique identifier base for tool anchor nodes. Tool nodes use IDs above any
/// realistic token/sense space so they never collide with vocabulary nodes.
const TOOL_ANCHOR_BASE: u64 = 10_000_000_000;
const TOOL_ARG_BASE: u64 = 20_000_000_000;

/// Properties tag for a tool anchor node.
pub const PROP_TOOL_ANCHOR: &str = "tool_anchor";
/// Properties tag for a tool argument node.
pub const PROP_TOOL_ARG: &str = "tool_arg";

/// Inject tool schema subgraphs into an existing graph.
///
/// For each tool:
///   - create an anchor node tagged with the schema,
///   - create argument nodes for each parameter,
///   - link anchor → argument with high-weight edges.
///
/// Anchor positions are placed at the centroid of the token sense nodes that
/// spell the tool name, falling back to a position in the tool domain region if
/// no senses are found.
pub fn inject_tool_subgraphs(
    graph: &mut Vec<GraphNode4D>,
    schemas: &[ToolSchema],
    tokenizer: &tokenizers::Tokenizer,
    domain_centroid: Option<Vec3>,
) -> Result<()> {
    let node_index = graph
        .iter()
        .enumerate()
        .map(|(i, n)| (n.id, i))
        .collect::<HashMap<_, _>>();

    let mut next_anchor_id = TOOL_ANCHOR_BASE;
    let mut next_arg_id = TOOL_ARG_BASE;

    let fallback_centroid = domain_centroid.unwrap_or_else(|| {
        // Centroid of existing graph nodes.
        let mut sum = Vec3::ZERO;
        for node in graph.iter() {
            sum += node.position();
        }
        sum / graph.len().max(1) as f32
    });

    let n_schemas = schemas.len().max(1);
    // Spread tools that fall back to the domain centroid across a sphere so they
    // do not collapse to a single point.
    let fallback_radius = (n_schemas as f32).cbrt() * 0.3;

    for (idx, schema) in schemas.iter().enumerate() {
        // Find sense nodes for tool name tokens.
        let anchor_pos = centroid_for_text(graph, &node_index, tokenizer, &schema.name)
            .unwrap_or_else(|| {
                fibonacci_sphere_point(fallback_centroid, fallback_radius, idx, n_schemas)
            });

        let anchor_id = next_anchor_id;
        next_anchor_id += 1;

        let mut props = GraphProperties::new();
        props.insert(
            "type".to_string(),
            serde_json::Value::String(PROP_TOOL_ANCHOR.to_string()),
        );
        props.insert(
            "name".to_string(),
            serde_json::Value::String(schema.name.clone()),
        );
        props.insert(
            "schema".to_string(),
            serde_json::to_value(schema).context("failed to serialize tool schema")?,
        );

        let anchor = GraphNode4D {
            id: anchor_id,
            x: anchor_pos.x,
            y: anchor_pos.y,
            z: anchor_pos.z,
            begin_ts: 0,
            end_ts: u64::MAX,
            properties: props,
            successors: Vec::new(),
        };
        graph.push(anchor);
        let anchor_graph_idx = graph.len() - 1;

        // Argument nodes arranged radially around the anchor.
        let params: Vec<String> = schema.required_params();
        let n = params.len().max(1) as f32;
        for (i, param) in params.iter().enumerate() {
            let angle = 2.0 * std::f32::consts::PI * (i as f32) / n;
            let radius = 0.5;
            let offset = Vec3::new(angle.cos() * radius, angle.sin() * radius, 0.0);
            let arg_pos = anchor_pos + offset;

            let arg_id = next_arg_id;
            next_arg_id += 1;

            let mut arg_props = GraphProperties::new();
            arg_props.insert(
                "type".to_string(),
                serde_json::Value::String(PROP_TOOL_ARG.to_string()),
            );
            arg_props.insert(
                "tool".to_string(),
                serde_json::Value::String(schema.name.clone()),
            );
            arg_props.insert("arg".to_string(), serde_json::Value::String(param.clone()));

            let arg = GraphNode4D {
                id: arg_id,
                x: arg_pos.x,
                y: arg_pos.y,
                z: arg_pos.z,
                begin_ts: 0,
                end_ts: u64::MAX,
                properties: arg_props,
                successors: Vec::new(),
            };
            graph.push(arg);

            // Anchor -> argument edge.
            graph[anchor_graph_idx].successors.push(TemporalEdge {
                dst: arg_id,
                weight: 1.0,
                begin_ts: 0,
                end_ts: u64::MAX,
            });
        }
    }

    Ok(())
}

/// Deterministic point on a sphere using a Fibonacci lattice.
fn fibonacci_sphere_point(center: Vec3, radius: f32, index: usize, total: usize) -> Vec3 {
    let phi = (1.0 + 5.0f32.sqrt()) / 2.0;
    let y = 1.0 - (2.0 * index as f32 + 1.0) / total.max(1) as f32;
    let theta = 2.0 * std::f32::consts::PI * index as f32 / phi;
    let r = (1.0 - y * y).sqrt();
    let x = r * theta.cos();
    let z = r * theta.sin();
    center + Vec3::new(x, y, z) * radius
}

/// Compute the centroid of sense nodes for all tokens in `text`.
fn centroid_for_text(
    graph: &[GraphNode4D],
    node_index: &HashMap<u64, usize>,
    tokenizer: &tokenizers::Tokenizer,
    text: &str,
) -> Option<Vec3> {
    let encoding = tokenizer.encode(text.to_string(), false).ok()?;
    let token_ids: Vec<u32> = encoding.get_ids().to_vec();
    if token_ids.is_empty() {
        return None;
    }

    // Collect sense nodes for each token ID.
    let mut points = Vec::new();
    for &tid in &token_ids {
        let base = (tid as u64) * 1000;
        for sense_offset in 0..1000 {
            let node_id = base + sense_offset;
            if let Some(&idx) = node_index.get(&node_id) {
                points.push(graph[idx].position());
            }
        }
    }

    if points.is_empty() {
        return None;
    }

    let sum: Vec3 = points.iter().sum();
    Some(sum / points.len() as f32)
}

/// Load tool schemas that were persisted alongside a graph.
///
/// Properties may be stringified during binary persistence, so this function
/// handles both `Value::Object` and `Value::String(JSON)` forms.
pub fn load_persisted_tool_schemas(graph: &[GraphNode4D]) -> Vec<ToolSchema> {
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    for node in graph {
        if node.properties.get("type").and_then(|v| v.as_str()) != Some(PROP_TOOL_ANCHOR) {
            continue;
        }
        let Some(name) = node.properties.get("name").and_then(|v| v.as_str()) else {
            continue;
        };
        if seen.contains(name) {
            continue;
        }
        let schema = node.properties.get("schema").and_then(|v| {
            let json = match v {
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string(),
            };
            serde_json::from_str::<ToolSchema>(&json).ok()
        });
        if let Some(schema) = schema {
            seen.insert(name.to_string());
            out.push(schema);
        }
    }
    out
}

/// Find the tool anchor node closest to a position, if within `radius`.
pub fn nearest_tool_anchor(
    graph: &[GraphNode4D],
    position: Vec3,
    radius: f32,
) -> Option<(usize, ToolSchema)> {
    nearest_tool_anchor_scored(graph, position, radius, &[], None)
}

/// Find a tool anchor near `position`, ranking candidates by token overlap with
/// the prompt. Falls back to pure spatial proximity if no overlap is found.
pub fn nearest_tool_anchor_for_prompt(
    graph: &[GraphNode4D],
    position: Vec3,
    radius: f32,
    prompt_tokens: &[u32],
    tokenizer: &tokenizers::Tokenizer,
) -> Option<(usize, ToolSchema)> {
    nearest_tool_anchor_scored(graph, position, radius, prompt_tokens, Some(tokenizer))
}

fn nearest_tool_anchor_scored(
    graph: &[GraphNode4D],
    position: Vec3,
    radius: f32,
    prompt_tokens: &[u32],
    tokenizer: Option<&tokenizers::Tokenizer>,
) -> Option<(usize, ToolSchema)> {
    let prompt_set: std::collections::HashSet<u32> = prompt_tokens.iter().copied().collect();

    let mut best = None;
    let mut best_score = 0i32;
    let mut best_dist = f32::MAX;

    for (i, node) in graph.iter().enumerate() {
        if node.properties.get("type").and_then(|v| v.as_str()) != Some(PROP_TOOL_ANCHOR) {
            continue;
        }
        let dist = node.position().distance(position);
        if dist > radius {
            continue;
        }

        let name = node
            .properties
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        let overlap = tokenizer
            .and_then(|t| t.encode(name.to_string(), false).ok())
            .map(|enc| {
                enc.get_ids()
                    .iter()
                    .filter(|tid| prompt_set.contains(tid))
                    .count() as i32
            })
            .unwrap_or(0);

        // Prefer overlap, break ties by distance.
        let improve = match best {
            None => true,
            Some(_) if overlap > best_score => true,
            Some(_) if overlap == best_score && dist < best_dist => true,
            _ => false,
        };

        if improve {
            best = Some((i, node));
            best_score = overlap;
            best_dist = dist;
        }
    }

    let (idx, node) = best?;
    node.properties
        .get("schema")
        .and_then(|v| {
            let json = match v {
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string(),
            };
            serde_json::from_str::<ToolSchema>(&json).ok()
        })
        .map(|schema| (idx, schema))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::four_d::GraphProperties;

    fn make_node(id: u64, x: f32, y: f32, z: f32) -> GraphNode4D {
        GraphNode4D {
            id,
            x,
            y,
            z,
            begin_ts: 0,
            end_ts: u64::MAX,
            properties: GraphProperties::new(),
            successors: Vec::new(),
        }
    }

    #[test]
    fn extract_json_object_finds_nested_object() {
        let text = r#"SYSTEM: you have functions -
{"name":"get_rate","description":"x","parameters":{"type":"object","properties":{"base":{"type":"string"}},"required":["base"]}}
more text"#;
        let json = extract_json_object(text).unwrap();
        assert!(json.contains("\"name\":\"get_rate\""));
        let schema: ToolSchema = serde_json::from_str(&json).unwrap();
        assert_eq!(schema.name, "get_rate");
        assert_eq!(schema.required_params(), vec!["base"]);
    }

    #[test]
    fn parse_tool_schemas_deduplicates() {
        let texts = vec![
            "{\"name\":\"a\",\"parameters\":{}}".to_string(),
            "{\"name\":\"a\",\"parameters\":{}}".to_string(),
            "{\"name\":\"b\",\"parameters\":{}}".to_string(),
        ];
        let schemas = parse_tool_schemas(&texts);
        assert_eq!(schemas.len(), 2);
    }

    #[test]
    fn load_persisted_tool_schemas_roundtrips() {
        let schema = ToolSchema {
            name: "foo".to_string(),
            description: "bar".to_string(),
            parameters: serde_json::json!({"required": ["x"]}),
        };
        let mut props = GraphProperties::new();
        props.insert(
            "type".to_string(),
            serde_json::Value::String(PROP_TOOL_ANCHOR.to_string()),
        );
        props.insert(
            "name".to_string(),
            serde_json::Value::String("foo".to_string()),
        );
        props.insert("schema".to_string(), serde_json::to_value(&schema).unwrap());
        let node = GraphNode4D {
            id: TOOL_ANCHOR_BASE,
            x: 0.0,
            y: 0.0,
            z: 0.0,
            begin_ts: 0,
            end_ts: u64::MAX,
            properties: props,
            successors: Vec::new(),
        };
        let loaded = load_persisted_tool_schemas(&[node]);
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded[0].name, "foo");
    }

    #[test]
    fn tool_anchor_persists_through_save_load() {
        let dir = tempfile::tempdir().unwrap();
        let mut graph = vec![make_node(0, 0.0, 0.0, 0.0), make_node(1, 1.0, 0.0, 0.0)];
        let schema = ToolSchema {
            name: "get_rate".to_string(),
            description: "".to_string(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {"base": {"type": "string"}},
                "required": ["base"]
            }),
        };

        // Build a fake tokenizer that maps "get_rate" to token id 7.
        let mut tokenizer = tokenizers::Tokenizer::new(tokenizers::models::bpe::BPE::default());
        tokenizer.add_tokens(&[tokenizers::AddedToken::from("get_rate".to_string(), false)]);

        inject_tool_subgraphs(&mut graph, &[schema], &tokenizer, None).unwrap();
        crate::save_graph4d(&graph, dir.path()).unwrap();
        let loaded = crate::load_graph4d(dir.path()).unwrap();

        let schemas = load_persisted_tool_schemas(&loaded);
        assert_eq!(schemas.len(), 1);
        assert_eq!(schemas[0].name, "get_rate");
        assert_eq!(schemas[0].required_params(), vec!["base"]);

        let anchor = nearest_tool_anchor(&loaded, Vec3::new(0.0, 0.0, 0.0), 10.0);
        assert!(anchor.is_some());
    }
}