needle_lib 0.1.1

//! # Needle Lib
//!
//! `needle_lib` is a high-performance Rust inference library for the Cactus Needle2 binary model.
//! It supports local tool-calling, LLM reasoning, confidence scoring, and structured output.
//!
//! ## Features
//! - **Local Inference**: Load precompiled weights and perform inference locally.
//! - **Tool Calling**: Define available tools in JSON schema format, and the engine automatically triggers them when needed.
//! - **Reasoning & Validation**: Retrieve reasoning steps and ungrounded validation errors from the model response.
//! - **Performance Metrics**: Access prefill TPS, decode TPS, and peak RAM consumption.
//!
//! ## Example
//!
//! ```no_run
//! use needle_lib::{Needle, ResponseType};
//! use serde_json::json;
//!
//! fn main() -> anyhow::Result<()> {
//!     // 1. Optionally load weights (if required by your environment setup)
//!     // Needle::load_weights(&weights_bytes)?;
//!
//!     // 2. Define tools in JSON schema format
//!     let tools = json!([{
//!         "name": "get_weather",
//!         "description": "Get the current weather for a city.",
//!         "parameters": {
//!             "type": "object",
//!             "properties": { "city": { "type": "string" } },
//!             "required": ["city"]
//!         }
//!     }]);
//!
//!     // 3. Initialize the Needle engine
//!     let needle = Needle::init("You are a helpful assistant.", &tools.to_string())?;
//!
//!     // 4. Run inference
//!     let response = needle.complete("What's the weather like in Paris?", 256)?;
//!
//!     if response.kind == ResponseType::Call {
//!         for call in &response.function_calls {
//!             println!("Tool called: {}", call.name);
//!             println!("Arguments: {:?}", call.arguments);
//!         }
//!     }
//!
//!     println!("Confidence: {}", response.confidence);
//!     println!("Reasoning: {}", response.reasoning.as_deref().unwrap_or("None"));
//!     Ok(())
//! }
//! ```

#![allow(non_upper_case_globals, non_camel_case_types, non_snake_case, dead_code)]

mod ffi {
    include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
}
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::os::raw::c_char;

/// The type of response returned by the Needle engine.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ResponseType {
    /// The model decided to invoke one or more external tools/functions.
    Call,
    /// The model finished reasoning and provided a final text response.
    Respond,
}

/// Represents the structured output from the Needle engine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeedleResponse {
    /// The type of response: either a tool call or a final response.
    #[serde(rename = "type")]
    pub kind: ResponseType,

    /// Whether the inference request was successfully processed.
    pub success: bool,

    /// An optional error message if the inference failed.
    pub error: Option<String>,

    /// An optional error code if the inference failed.
    pub error_code: Option<String>,

    /// The optional reason behind a failure or specific decision.
    pub reason: Option<String>,

    /// A list of function calls triggered by the model.
    pub function_calls: Vec<FunctionCall>,

    /// The model's inner reasoning/thinking steps, if enabled and generated.
    pub reasoning: Option<String>,

    /// A confidence score indicating how certain the model is about its output.
    pub confidence: f64,

    /// Optional grounding validation output.
    pub validation: Option<Validation>,

    /// Token-per-second rate during the prefill phase.
    pub prefill_tps: f64,

    /// Token-per-second rate during the decode phase.
    pub decode_tps: f64,

    /// Peak RAM usage in Megabytes during the inference.
    pub peak_ram_mb: f64,
}

/// Represents a single function call generated by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    /// The name of the function to be called.
    pub name: String,

    /// The structured arguments passed to the function, mapped by parameter name.
    pub arguments: HashMap<String, serde_json::Value>,
}

/// Grounding validation details for the generated response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Validation {
    /// Whether negation was detected in the validation step.
    pub negation: bool,

    /// A list of ungrounded statements or details detected during validation.
    pub ungrounded: Vec<String>,
}

/// The core Needle engine interface, wrapping the native pre-compiled C-library.
pub struct Needle;

impl Needle {
    /// Initializes the Needle engine with a system prompt and a set of JSON-defined tools.
    ///
    /// # Arguments
    /// * `system` - A system prompt defining the instructions or persona of the assistant.
    /// * `tools_json` - A JSON-string representing an array of available tools/functions.
    ///
    /// # Errors
    /// Returns an error if the initialization of the underlying native C engine fails.
    pub fn init(system: &str, tools_json: &str) -> anyhow::Result<Self> {
        let system_c = CString::new(system)?;
        let tools_c = CString::new(tools_json)?;
        let ret = unsafe {
            ffi::needle_init(system_c.as_ptr(), tools_c.as_ptr(), std::ptr::null())
        };
        if ret < 0 {
            anyhow::bail!("needle_init failed with code {ret}");
        }
        Ok(Needle)
    }

    /// Completes a text prompt by executing local inference.
    ///
    /// # Arguments
    /// * `text` - The user prompt or text input for the engine.
    /// * `max_new_tokens` - The maximum number of tokens to generate.
    ///
    /// # Errors
    /// Returns an error if the native engine fails to complete the text, or if the resulting
    /// buffer cannot be parsed into valid UTF-8 and deserialized to a `NeedleResponse`.
    pub fn complete(&self, text: &str, max_new_tokens: i32) -> anyhow::Result<NeedleResponse> {
        let text_c = CString::new(text)?;
        let mut buf = vec![0u8; 65536];

        unsafe {
            ffi::needle_complete(
                text_c.as_ptr(),
                max_new_tokens,
                buf.as_mut_ptr() as *mut c_char,
                buf.len() as i32,
            );
        }

        let cstr = unsafe { CStr::from_ptr(buf.as_ptr() as *const c_char) };
        let cstr_str = cstr.to_str()?;
        Ok(serde_json::from_str(cstr_str)?)
    }

    /// Resets the internal state/context of the native Needle engine.
    pub fn reset(&self) {
        unsafe { ffi::needle_reset() };
    }

    /// Loads the binary weights blob into the native Needle engine.
    ///
    /// # Arguments
    /// * `blob` - A byte slice containing the compiled weights of the Cactus Needle2 model.
    ///
    /// # Errors
    /// Returns an error if the binary load fails in the underlying native engine.
    pub fn load_weights(blob: &[u8]) -> anyhow::Result<()> {
        let ret = unsafe { ffi::needle_load(blob.as_ptr() as *const c_char, blob.len() as u64) };
        if ret != 0 {
            anyhow::bail!("needle_load failed with code {ret}");
        }
        Ok(())
    }
}