cortexdb-sdk 0.2.0-beta.2

Rust HTTP client SDK for CortexDB
Documentation
use crate::{
    AnswerGroundingOptionsResponse, AqlRetrievalMode, ContextPackCellResponse, ContextPackResponse,
    GroundedAnswerRequest, GroundedAnswerResponse,
};

fn test_context_pack() -> ContextPackResponse {
    ContextPackResponse {
        schema_version: ContextPackResponse::SCHEMA_VERSION_V1.to_owned(),
        token_budget_tokens: 256,
        estimated_tokens: 40,
        truncated: false,
        citations_required: true,
        answerability_q16: u16::MAX,
        conflict_visibility_q16: 0,
        visible_conflict_count: 0,
        cells: vec![ContextPackCellResponse {
            cell_id: 7,
            estimated_tokens: 40,
            citation: Some("doc://project-risk#p1".to_owned()),
            payload_text: "The migration blocker is the audit export dependency.".to_owned(),
            explain: None,
            source_ref: None,
            provenance: None,
            access_decision: None,
        }],
        anomalies: Vec::new(),
    }
}

#[test]
fn grounded_answer_request_builds_retrieve_and_verify_statements() {
    let request = GroundedAnswerRequest::new("project:alpha", "default", "migration blocker")
        .mode(AqlRetrievalMode::Balanced)
        .budget_tokens(512)
        .limit_candidates(12)
        .where_clause(r#"space = project:alpha AND status = "ready""#)
        .require_citations(true);

    assert_eq!(
        request.retrieve_statement().unwrap(),
        concat!(
            "RETRIEVE CONTEXT FOR TASK \"migration blocker\" IN BRAIN default ",
            "USING MODE balanced BUDGET 512 TOKENS LIMIT 12 CANDIDATES ",
            "WHERE space = project:alpha AND status = \"ready\" REQUIRE citations;"
        )
    );
    assert_eq!(
        request.verify_statement("The migration blocker is the audit export dependency.").unwrap(),
        Some(
            "VERIFY FACT \"The migration blocker is the audit export dependency.\" IN BRAIN default;"
                .to_owned()
        )
    );
}

#[test]
fn grounded_answer_response_collects_citations_and_grounding() {
    let request = GroundedAnswerRequest::new("project:alpha", "default", "migration blocker")
        .grounding_options(AnswerGroundingOptionsResponse {
            require_citations: true,
            reject_unsupported: true,
            ..AnswerGroundingOptionsResponse::default()
        });
    let response = GroundedAnswerResponse::from_context_answer(
        &request,
        request.retrieve_statement().unwrap(),
        request
            .verify_statement("The migration blocker is the audit export dependency.")
            .unwrap(),
        test_context_pack(),
        "The migration blocker is the audit export dependency.".to_owned(),
        None,
    );

    assert!(response.answer_supported());
    assert!(!response.rejected);
    assert_eq!(response.citations, vec!["doc://project-risk#p1"]);
    assert_eq!(response.used_context_cell_ids, vec![7]);
    assert_eq!(
        response.verify_statement.as_deref(),
        Some(
            "VERIFY FACT \"The migration blocker is the audit export dependency.\" IN BRAIN default;"
        )
    );
}

#[test]
fn grounded_answer_can_disable_verify_for_draft_answers() {
    let request = GroundedAnswerRequest::new("project:alpha", "default", "migration blocker")
        .verify_answer(false);

    assert_eq!(
        request
            .verify_statement("The migration blocker is the audit export dependency.")
            .unwrap(),
        None
    );
}