kkachi 0.1.8

High-performance, zero-copy library for optimizing language model prompts and programs
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

// Copyright © 2025 lituus-io <spicyzhug@gmail.com>
// All Rights Reserved.
// Licensed under PolyForm Noncommercial 1.0.0

//! Integration tests for Predict module

use kkachi::buffer::Buffer;
use kkachi::predict::{predict_with_lm, LMClient, LMOutput, Predict};
use kkachi::str_view::StrView;
use kkachi::{Result, Signature};

/// Test LM client that returns predefined responses
struct TestLM {
    response: &'static str,
}

impl TestLM {
    fn new(response: &'static str) -> Self {
        Self { response }
    }
}

impl LMClient for TestLM {
    type GenerateFut<'a>
        = std::future::Ready<Result<LMOutput<'a>>>
    where
        Self: 'a;

    fn generate<'a>(&'a self, _prompt: StrView<'a>) -> Self::GenerateFut<'a> {
        // Create a buffer with the response - using static for simplicity
        // In real usage, this would be from the HTTP response
        static RESPONSE_1: Buffer = Buffer::Static(b"Answer: Paris");
        static RESPONSE_2: Buffer = Buffer::Static(b"Answer: 8");
        static RESPONSE_3: Buffer = Buffer::Static(b"Sentiment: positive\nConfidence: 0.95");

        // Select response based on content
        let buffer = if self.response.contains("Paris") {
            &RESPONSE_1
        } else if self.response.contains("8") {
            &RESPONSE_2
        } else {
            &RESPONSE_3
        };

        std::future::ready(Ok(LMOutput {
            buffer: buffer.view_all(),
            prompt_tokens: 10,
            completion_tokens: 5,
        }))
    }
}

#[tokio::test]
async fn test_predict_basic_qa() {
    let sig = Signature::parse("question -> answer").unwrap();
    let predict = Predict::without_demos(&sig);
    let lm = TestLM::new("Paris");

    let mut inputs = kkachi::types::Inputs::new();
    inputs.insert("question", "What is the capital of France?");

    let mut prompt_buffer = Vec::new();
    let result = predict_with_lm(&predict, &inputs, &lm, &mut prompt_buffer).await;

    assert!(result.is_ok());
    let output = result.unwrap();
    let answer = output.get_by_name("answer");
    assert!(answer.is_some());
    assert!(answer.unwrap().as_str().contains("Paris"));
}

#[tokio::test]
async fn test_predict_multiple_outputs() {
    let sig = Signature::parse("text -> sentiment, confidence").unwrap();
    let predict = Predict::without_demos(&sig);
    let lm = TestLM::new("sentiment");

    let mut inputs = kkachi::types::Inputs::new();
    inputs.insert("text", "This is great!");

    let mut prompt_buffer = Vec::new();
    let result = predict_with_lm(&predict, &inputs, &lm, &mut prompt_buffer).await;

    assert!(result.is_ok());
    let output = result.unwrap();

    // Note: parsing depends on field prefixes - this test validates the flow works
    assert!(output.prompt_tokens > 0);
    assert!(output.completion_tokens > 0);
}

#[tokio::test]
async fn test_predict_builds_prompt() {
    let sig = Signature::parse("question -> answer").unwrap();
    let predict = Predict::without_demos(&sig);

    let mut inputs = kkachi::types::Inputs::new();
    inputs.insert("question", "Test question?");

    let mut buffer = Vec::new();
    let prompt = predict.build_prompt_into(&inputs, &mut buffer);

    assert!(prompt.as_str().contains("Question"));
    assert!(prompt.as_str().contains("Test question?"));
    assert!(prompt.as_str().contains("Answer"));
}

#[tokio::test]
async fn test_predict_parse_response() {
    let sig = Signature::parse("question -> answer").unwrap();
    let predict = Predict::without_demos(&sig);

    // Test parsing a response
    let response = StrView::new("Answer: The result is 42");
    let ranges = predict.parse_response_ranges(response);

    assert!(!ranges.is_empty());
}