coyote/src/graph/structured.rs

use crate::client::call_chat_completions;
use crate::config::{Input, RequestContext, Role, RoleLike};
use crate::utils::create_abort_signal;
use anyhow::{Context, Result, bail};
use serde_json::Value;
use std::sync::Arc;

const EXTRACTOR_ROLE_NAME: &str = "__structured_output__";

const EXTRACTOR_ROLE_PROMPT: &str = "\
Extract a JSON object from the user's input that strictly conforms to the provided JSON Schema.

Rules:
- Output ONLY the JSON object. No prose, no explanation, no markdown fences, no <think> tokens.
- The first character of your response must be `{` and the last must be `}`.
- Every key marked `required` in the schema MUST appear in the output.
- All values MUST match the types specified in the schema.
- If the input is already a valid JSON object matching the schema, return it unchanged.
- If a field cannot be determined from the input, use `null` (when allowed) or your best inferred value.
- Do NOT invent fields not present in the schema.";

pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result<Value> {
    if let Some(parsed) = try_parse_json(raw) {
        return Ok(parsed);
    }

    extract_via_extractor(raw, schema, parent_ctx, false).await
}

async fn extract_via_extractor(
    raw: &str,
    schema: &Value,
    parent_ctx: &mut RequestContext,
    is_repair: bool,
) -> Result<Value> {
    let role = build_extractor_role()?;
    let prompt = build_extractor_prompt(raw, schema, is_repair);

    let saved_role = parent_ctx.role.clone();
    parent_ctx.role = Some(role);
    let result = run_one_shot(&prompt, parent_ctx).await;
    parent_ctx.role = saved_role;

    let output = result.context("Structured-output extractor LLM call failed")?;

    match try_parse_json(&output) {
        Some(value) => Ok(value),
        None if is_repair => bail!(
            "Structured-output extractor failed to produce valid JSON after repair retry. \
             Last response:\n{output}"
        ),
        None => Box::pin(extract_via_extractor(&output, schema, parent_ctx, true)).await,
    }
}

fn build_extractor_role() -> Result<Role> {
    let mut role = Role::new(EXTRACTOR_ROLE_NAME, EXTRACTOR_ROLE_PROMPT);
    role.set_enabled_tools(Some(Vec::new()));
    role.set_enabled_mcp_servers(Some(Vec::new()));
    Ok(role)
}

fn build_extractor_prompt(raw: &str, schema: &Value, is_repair: bool) -> String {
    let schema_json = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
    if is_repair {
        format!(
            "Your previous response was not valid JSON. Output ONLY a JSON object \
             matching this schema. No prose, no fences.\n\nSchema:\n{schema_json}\n\nInput:\n{raw}"
        )
    } else {
        format!("Schema:\n{schema_json}\n\nInput:\n{raw}")
    }
}

async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result<String> {
    let abort = create_abort_signal();
    let app_cfg = Arc::clone(&ctx.app.config);
    let role_for_input = ctx.role.clone();
    let input = Input::from_str(ctx, prompt, role_for_input)?;
    let client = input.create_client()?;
    ctx.before_chat_completion(&input)?;
    let (output, tool_results) =
        call_chat_completions(&input, false, false, client.as_ref(), ctx, abort).await?;
    ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;

    Ok(output)
}

fn try_parse_json(raw: &str) -> Option<Value> {
    let cleaned = strip_code_fences(raw.trim());

    serde_json::from_str(cleaned).ok()
}

fn strip_code_fences(s: &str) -> &str {
    let after_open = s
        .strip_prefix("```json")
        .or_else(|| s.strip_prefix("```"))
        .map(str::trim_start)
        .unwrap_or(s);
    after_open
        .strip_suffix("```")
        .map(str::trim_end)
        .unwrap_or(after_open)
}

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

    #[test]
    fn try_parse_json_accepts_plain_object() {
        let v = try_parse_json(r#"{"a": 1}"#).unwrap();

        assert_eq!(v, json!({"a": 1}));
    }

    #[test]
    fn try_parse_json_strips_json_fences() {
        let raw = "```json\n{\"a\": 1}\n```";

        let v = try_parse_json(raw).unwrap();

        assert_eq!(v, json!({"a": 1}));
    }

    #[test]
    fn try_parse_json_strips_bare_fences() {
        let raw = "```\n{\"a\": 1}\n```";

        let v = try_parse_json(raw).unwrap();

        assert_eq!(v, json!({"a": 1}));
    }

    #[test]
    fn try_parse_json_tolerates_whitespace() {
        let v = try_parse_json("   \n  {\"x\": true}\n\n").unwrap();

        assert_eq!(v, json!({"x": true}));
    }

    #[test]
    fn try_parse_json_returns_none_on_prose() {
        assert!(try_parse_json("Here is the result: it's good").is_none());
    }

    #[test]
    fn try_parse_json_returns_none_on_partial_json() {
        assert!(try_parse_json("{\"a\": ").is_none());
    }

    #[test]
    fn try_parse_json_accepts_arrays() {
        let v = try_parse_json("[1, 2, 3]").unwrap();

        assert_eq!(v, json!([1, 2, 3]));
    }

    #[test]
    fn build_extractor_prompt_includes_schema_and_input() {
        let schema = json!({"type": "object"});

        let prompt = build_extractor_prompt("hello", &schema, false);

        assert!(prompt.contains("Schema:"));
        assert!(prompt.contains("Input:"));
        assert!(prompt.contains("hello"));
    }

    #[test]
    fn build_extractor_prompt_repair_includes_repair_instruction() {
        let schema = json!({"type": "object"});

        let prompt = build_extractor_prompt("oops", &schema, true);

        assert!(prompt.contains("previous response"));
        assert!(prompt.contains("oops"));
    }

    #[test]
    fn build_extractor_role_disables_tools_and_mcp() {
        let role = build_extractor_role().expect("builtin role must exist");

        assert_eq!(role.enabled_tools().as_deref(), Some([].as_slice()));
        assert_eq!(role.enabled_mcp_servers().as_deref(), Some([].as_slice()));
    }
}