consciousness/src/agent/tools/memory.rs

// tools/memory.rs — Native memory graph operations
//
// Direct library calls into the store — no subprocess spawning.
// Returns MemoryNodes where possible so the agent can track what's
// loaded in its context window.

use anyhow::{Context, Result};
use serde_json::json;

use crate::agent::memory::MemoryNode;
use crate::agent::types::ToolDef;
use crate::store::{self, Store};

pub fn definitions() -> Vec<ToolDef> {
    vec![
        ToolDef::new(
            "memory_render",
            "Read a memory node's content and links. Returns the full content \
             with neighbor links sorted by strength.",
            json!({
                "type": "object",
                "properties": {
                    "key": {
                        "type": "string",
                        "description": "Node key to render"
                    }
                },
                "required": ["key"]
            }),
        ),
        ToolDef::new(
            "memory_write",
            "Create or update a memory node with new content. Use for writing \
             prose, analysis, or any node content. Multi-line content is fine.",
            json!({
                "type": "object",
                "properties": {
                    "key": {
                        "type": "string",
                        "description": "Node key to create or update"
                    },
                    "content": {
                        "type": "string",
                        "description": "Full content for the node (markdown)"
                    }
                },
                "required": ["key", "content"]
            }),
        ),
        ToolDef::new(
            "memory_search",
            "Search the memory graph for nodes by keyword.",
            json!({
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search terms"
                    }
                },
                "required": ["query"]
            }),
        ),
        ToolDef::new(
            "memory_links",
            "Show a node's neighbors with link strengths.",
            json!({
                "type": "object",
                "properties": {
                    "key": {
                        "type": "string",
                        "description": "Node key to show links for"
                    }
                },
                "required": ["key"]
            }),
        ),
        ToolDef::new(
            "memory_link_set",
            "Set the strength of a link between two nodes. Also deduplicates \
             if multiple links exist between the same pair.",
            json!({
                "type": "object",
                "properties": {
                    "source": {
                        "type": "string",
                        "description": "Source node key"
                    },
                    "target": {
                        "type": "string",
                        "description": "Target node key"
                    },
                    "strength": {
                        "type": "number",
                        "description": "Link strength (0.01 to 1.0)"
                    }
                },
                "required": ["source", "target", "strength"]
            }),
        ),
        ToolDef::new(
            "memory_link_add",
            "Add a new link between two nodes.",
            json!({
                "type": "object",
                "properties": {
                    "source": {
                        "type": "string",
                        "description": "Source node key"
                    },
                    "target": {
                        "type": "string",
                        "description": "Target node key"
                    }
                },
                "required": ["source", "target"]
            }),
        ),
        ToolDef::new(
            "memory_used",
            "Mark a node as useful (boosts its weight in the graph).",
            json!({
                "type": "object",
                "properties": {
                    "key": {
                        "type": "string",
                        "description": "Node key to mark as used"
                    }
                },
                "required": ["key"]
            }),
        ),
        ToolDef::new(
            "memory_weight_set",
            "Set a node's weight directly. Use to downweight junk nodes (0.01) \
             or boost important ones. Normal range is 0.1 to 1.0.",
            json!({
                "type": "object",
                "properties": {
                    "key": {
                        "type": "string",
                        "description": "Node key"
                    },
                    "weight": {
                        "type": "number",
                        "description": "New weight (0.01 to 1.0)"
                    }
                },
                "required": ["key", "weight"]
            }),
        ),
        ToolDef::new(
            "memory_supersede",
            "Mark a node as superseded by another. Sets the old node's weight \
             to 0.01 and prepends a notice pointing to the replacement. Use \
             when merging duplicates or replacing junk with proper content.",
            json!({
                "type": "object",
                "properties": {
                    "old_key": {
                        "type": "string",
                        "description": "Node being superseded"
                    },
                    "new_key": {
                        "type": "string",
                        "description": "Replacement node"
                    },
                    "reason": {
                        "type": "string",
                        "description": "Why this node was superseded (e.g. 'merged into X', 'duplicate of Y')"
                    }
                },
                "required": ["old_key", "new_key"]
            }),
        ),
    ]
}

/// Dispatch a memory tool call. Direct library calls, no subprocesses.
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
    let prov = provenance.unwrap_or("manual");
    let result = match name {
        "memory_render" => {
            let key = get_str(args, "key")?;
            let node = MemoryNode::load(key)
                .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
            node.render()
        }
        "memory_write" => {
            let key = get_str(args, "key")?;
            let content = get_str(args, "content")?;
            let node = MemoryNode::write(key, content, Some(prov))
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            format!("wrote '{}' (v{})", node.key, node.version)
        }
        "memory_search" => {
            let query = get_str(args, "query")?;
            let results = MemoryNode::search(query)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            if results.is_empty() {
                "no results".to_string()
            } else {
                results.iter().map(|r| r.render()).collect::<Vec<_>>().join("\n")
            }
        }
        "memory_links" => {
            let key = get_str(args, "key")?;
            let node = MemoryNode::load(key)
                .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
            let mut out = format!("Neighbors of '{}':\n", key);
            for link in &node.links {
                out.push_str(&format!("  ({:.2}) {}{}\n",
                    link.strength, link.target,
                    if link.inline { " [inline]" } else { "" }));
            }
            out
        }
        "memory_link_set" => {
            let source = get_str(args, "source")?;
            let target = get_str(args, "target")?;
            let strength = get_f64(args, "strength")? as f32;
            link_set(source, target, strength)?
        }
        "memory_link_add" => {
            let source = get_str(args, "source")?;
            let target = get_str(args, "target")?;
            link_add(source, target, prov)?
        }
        "memory_used" => {
            let key = get_str(args, "key")?;
            MemoryNode::mark_used(key)
                .map_err(|e| anyhow::anyhow!("{}", e))?
        }
        "memory_weight_set" => {
            let key = get_str(args, "key")?;
            let weight = get_f64(args, "weight")? as f32;
            weight_set(key, weight)?
        }
        "memory_supersede" => supersede(args, prov)?,
        _ => anyhow::bail!("Unknown memory tool: {}", name),
    };
    Ok(result)
}

fn link_set(source: &str, target: &str, strength: f32) -> Result<String> {
    let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
    let source = store.resolve_key(source).map_err(|e| anyhow::anyhow!("{}", e))?;
    let target = store.resolve_key(target).map_err(|e| anyhow::anyhow!("{}", e))?;
    let strength = strength.clamp(0.01, 1.0);

    let mut found = false;
    let mut first = true;
    for rel in &mut store.relations {
        if rel.deleted { continue; }
        if (rel.source_key == source && rel.target_key == target)
            || (rel.source_key == target && rel.target_key == source)
        {
            if first {
                let old = rel.strength;
                rel.strength = strength;
                first = false;
                found = true;
                if (old - strength).abs() < 0.001 {
                    return Ok(format!("{} ↔ {} already at {:.2}", source, target, strength));
                }
            } else {
                rel.deleted = true; // deduplicate
            }
        }
    }

    if !found {
        anyhow::bail!("no link found between {} and {}", source, target);
    }

    store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
    Ok(format!("set {} ↔ {} strength to {:.2}", source, target, strength))
}

fn link_add(source: &str, target: &str, prov: &str) -> Result<String> {
    let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
    let source = store.resolve_key(source).map_err(|e| anyhow::anyhow!("{}", e))?;
    let target = store.resolve_key(target).map_err(|e| anyhow::anyhow!("{}", e))?;

    // Check for existing link
    let exists = store.relations.iter().any(|r|
        !r.deleted &&
        ((r.source_key == source && r.target_key == target) ||
         (r.source_key == target && r.target_key == source)));
    if exists {
        return Ok(format!("link already exists: {} ↔ {}", source, target));
    }

    let source_uuid = store.nodes.get(&source)
        .map(|n| n.uuid)
        .ok_or_else(|| anyhow::anyhow!("source not found: {}", source))?;
    let target_uuid = store.nodes.get(&target)
        .map(|n| n.uuid)
        .ok_or_else(|| anyhow::anyhow!("target not found: {}", target))?;

    // Compute initial strength from Jaccard similarity
    let graph = store.build_graph();
    let jaccard = graph.jaccard(&source, &target);
    let strength = (jaccard * 3.0).clamp(0.1, 1.0) as f32;

    let mut rel = store::new_relation(
        source_uuid, target_uuid,
        store::RelationType::Link, strength,
        &source, &target,
    );
    rel.provenance = prov.to_string();
    store.add_relation(rel).map_err(|e| anyhow::anyhow!("{}", e))?;
    store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
    Ok(format!("linked {} → {} (strength={:.2})", source, target, strength))
}

fn weight_set(key: &str, weight: f32) -> Result<String> {
    let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
    let resolved = store.resolve_key(key).map_err(|e| anyhow::anyhow!("{}", e))?;
    let weight = weight.clamp(0.01, 1.0);

    let node = store.nodes.get_mut(&resolved)
        .ok_or_else(|| anyhow::anyhow!("node not found: {}", resolved))?;
    let old = node.weight;
    node.weight = weight;
    store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
    Ok(format!("weight {} {:.2} → {:.2}", resolved, old, weight))
}

fn supersede(args: &serde_json::Value, prov: &str) -> Result<String> {
    let old_key = get_str(args, "old_key")?;
    let new_key = get_str(args, "new_key")?;
    let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");

    // Load old node
    let old = MemoryNode::load(old_key)
        .ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?;

    // Prepend superseded notice (strip link footer from content)
    let content_only = old.content.split("\n\n---\nLinks:").next().unwrap_or(&old.content);
    let notice = format!(
        "**SUPERSEDED** by `{}` — {}\n\nOriginal content preserved below for reference.\n\n---\n\n{}",
        new_key, reason, content_only.trim()
    );

    // Write back + set weight
    MemoryNode::write(old_key, &notice, Some(prov))
        .map_err(|e| anyhow::anyhow!("{}", e))?;
    weight_set(old_key, 0.01)?;

    Ok(format!("superseded {} → {} ({})", old_key, new_key, reason))
}

/// Helper: get required string argument.
fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
    args.get(name)
        .and_then(|v| v.as_str())
        .context(format!("{} is required", name))
}

/// Helper: get required f64 argument.
fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
    args.get(name)
        .and_then(|v| v.as_f64())
        .context(format!("{} is required", name))
}