mubit-sdk 0.7.0

//! MuBit Learn Context Injection.
//!
//! Functions for injecting context blocks into LLM message arrays
//! represented as `Vec<serde_json::Value>`.

use serde_json::{json, Value};

const MEMORY_TAG: &str = "<memory_context>";
const MEMORY_TAG_END: &str = "</memory_context>";

/// Inject a context block into an OpenAI-format message array.
///
/// Returns a new `Vec<Value>` with the context block injected
/// according to `position`: `"system"`, `"prepend"`, or `"last_system"`.
pub fn inject_context_openai(
    messages: &[Value],
    context_block: &str,
    position: &str,
) -> Vec<Value> {
    if context_block.trim().is_empty() {
        return messages.to_vec();
    }

    let mut msgs: Vec<Value> = messages.to_vec();
    let memory_text = format!("\n\n---\n{MEMORY_TAG}\n{context_block}\n{MEMORY_TAG_END}");

    match position {
        "system" => {
            if !msgs.is_empty() && msgs[0].get("role").and_then(|v| v.as_str()) == Some("system")
            {
                let existing = msgs[0]
                    .get("content")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                msgs[0] = json!({
                    "role": "system",
                    "content": format!("{existing}{memory_text}"),
                });
            } else {
                msgs.insert(0, json!({"role": "system", "content": memory_text.trim()}));
            }
        }
        "prepend" => {
            msgs.insert(0, json!({"role": "system", "content": memory_text.trim()}));
        }
        "last_system" => {
            let mut inserted = false;
            for i in (0..msgs.len()).rev() {
                if msgs[i].get("role").and_then(|v| v.as_str()) == Some("user") {
                    msgs.insert(
                        i,
                        json!({"role": "system", "content": memory_text.trim()}),
                    );
                    inserted = true;
                    break;
                }
            }
            if !inserted {
                msgs.insert(0, json!({"role": "system", "content": memory_text.trim()}));
            }
        }
        _ => {
            // Default to "system" behavior
            if !msgs.is_empty() && msgs[0].get("role").and_then(|v| v.as_str()) == Some("system")
            {
                let existing = msgs[0]
                    .get("content")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                msgs[0] = json!({
                    "role": "system",
                    "content": format!("{existing}{memory_text}"),
                });
            } else {
                msgs.insert(0, json!({"role": "system", "content": memory_text.trim()}));
            }
        }
    }

    msgs
}

/// Extract a query string from the last user message in the array.
pub fn extract_query(messages: &[Value], max_length: usize) -> String {
    for msg in messages.iter().rev() {
        if msg.get("role").and_then(|v| v.as_str()) == Some("user") {
            if let Some(content) = msg.get("content").and_then(|v| v.as_str()) {
                let end = content.len().min(max_length);
                return content[..end].to_string();
            }
        }
    }
    String::new()
}