coyote/src/graph/structured.rs

use crate::client::call_chat_completions_streaming_quiet;
use crate::config::{Input, RequestContext, Role, RoleLike};
use crate::utils::AbortSignal;
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,
    abort: &AbortSignal,
) -> Result<Value> {
    if let Some(parsed) = try_parse_json(raw) {
        return Ok(parsed);
    }

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

async fn extract_via_extractor(
    raw: &str,
    schema: &Value,
    parent_ctx: &mut RequestContext,
    abort: &AbortSignal,
    is_repair: bool,
) -> Result<Value> {
    let role = build_extractor_role(parent_ctx);
    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, abort).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, abort, true,
            ))
            .await
        }
    }
}

fn build_extractor_role(ctx: &RequestContext) -> Role {
    let mut role = Role::new(EXTRACTOR_ROLE_NAME, EXTRACTOR_ROLE_PROMPT);
    role.set_model(ctx.current_model().clone());
    role.set_enabled_tools(Some(Vec::new()));
    role.set_enabled_mcp_servers(Some(Vec::new()));
    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,
    abort: &AbortSignal,
) -> Result<String> {
    if abort.aborted() {
        bail!("Aborted.");
    }
    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_streaming_quiet(&input, client.as_ref(), ctx, abort.clone()).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(strip_thinking_blocks(raw.trim()));
    serde_json::from_str(cleaned).ok()
}

fn strip_thinking_blocks(s: &str) -> &str {
    let mut s = s.trim_start();
    while s.starts_with("<think>") {
        match s.find("</think>") {
            Some(end) => s = s[end + "</think>".len()..].trim_start(),
            None => break,
        }
    }
    s
}

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 crate::client::Model;
    use crate::config::{AppState, WorkingMode};
    use serde_json::json;

    fn make_ctx() -> RequestContext {
        RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd)
    }

    #[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_strips_thinking_blocks() {
        let raw = "<think>\nsome reasoning\n</think>\n{\"a\": 1}";

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

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

    #[test]
    fn try_parse_json_strips_empty_thinking_block() {
        let raw = "<think>\n\n</think>\n{\"a\": 1}";

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

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

    #[test]
    fn try_parse_json_strips_multiple_thinking_blocks() {
        let raw = "<think>first</think>\n<think>second</think>\n{\"a\": 1}";

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

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

    #[test]
    fn try_parse_json_unclosed_think_tag_returns_none() {
        assert!(try_parse_json("<think>unclosed {\"a\": 1}").is_none());
    }

    #[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 ctx = make_ctx();

        let role = build_extractor_role(&ctx);

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

    #[test]
    fn build_extractor_role_uses_parent_context_model() {
        let mut ctx = make_ctx();
        let mut parent_role = Role::new("parent", "parent prompt");
        parent_role.set_model(Model::new("client-x", "model-y"));
        ctx.role = Some(parent_role);

        let role = build_extractor_role(&ctx);

        assert_eq!(role.model().id(), "client-x:model-y");
    }
}