rskit-llm-openai 0.2.0-alpha.3

OpenAI provider implementation for rskit
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

use serde_json::json;

use rskit_llm_common as common;
use rskit_llm_common::OpenAiDialect;

struct FixtureCase<'a> {
    name: &'a str,
    body: &'a str,
    expected: Vec<rskit_ai::ToolUseBlock>,
}

fn streamed_tool_uses(events: &[&[u8]]) -> Vec<rskit_ai::ToolUseBlock> {
    let chunks = events
        .iter()
        .map(|event| OpenAiDialect::parse_stream_chunk(event).unwrap())
        .collect::<Vec<_>>();
    common::accumulate_tool_uses(chunks).unwrap()
}

#[test]
#[allow(clippy::too_many_lines)]
fn parses_non_streaming_fixtures() {
    let cases = [
        FixtureCase {
            name: "single tool call",
            body: r#"{
                "model": "gpt-4o",
                "choices": [{
                    "message": {
                        "content": null,
                        "tool_calls": [{
                            "id": "call_1",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": "{\"location\":\"NYC\"}"
                            }
                        }]
                    },
                    "finish_reason": "tool_calls"
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 5}
            }"#,
            expected: vec![rskit_ai::ToolUseBlock {
                id: "call_1".into(),
                name: "get_weather".into(),
                input: json!({"location": "NYC"}).as_object().cloned().unwrap(),
            }],
        },
        FixtureCase {
            name: "multi tool response",
            body: r#"{
                "model": "gpt-4o",
                "choices": [{
                    "message": {
                        "content": null,
                        "tool_calls": [
                            {
                                "id": "call_1",
                                "type": "function",
                                "function": {
                                    "name": "get_weather",
                                    "arguments": "{\"location\":\"NYC\"}"
                                }
                            },
                            {
                                "id": "call_2",
                                "type": "function",
                                "function": {
                                    "name": "get_time",
                                    "arguments": "{\"timezone\":\"UTC\"}"
                                }
                            }
                        ]
                    },
                    "finish_reason": "tool_calls"
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 5}
            }"#,
            expected: vec![
                rskit_ai::ToolUseBlock {
                    id: "call_1".into(),
                    name: "get_weather".into(),
                    input: json!({"location": "NYC"}).as_object().cloned().unwrap(),
                },
                rskit_ai::ToolUseBlock {
                    id: "call_2".into(),
                    name: "get_time".into(),
                    input: json!({"timezone": "UTC"}).as_object().cloned().unwrap(),
                },
            ],
        },
        FixtureCase {
            name: "empty args become empty object",
            body: r#"{
                "model": "gpt-4o",
                "choices": [{
                    "message": {
                        "content": null,
                        "tool_calls": [{
                            "id": "call_1",
                            "type": "function",
                            "function": {
                                "name": "ping",
                                "arguments": "{}"
                            }
                        }]
                    },
                    "finish_reason": "tool_calls"
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 5}
            }"#,
            expected: vec![rskit_ai::ToolUseBlock {
                id: "call_1".into(),
                name: "ping".into(),
                input: serde_json::Map::new(),
            }],
        },
        FixtureCase {
            name: "nested input",
            body: r#"{
                "model": "gpt-4o",
                "choices": [{
                    "message": {
                        "content": null,
                        "tool_calls": [{
                            "id": "call_1",
                            "type": "function",
                            "function": {
                                "name": "search",
                                "arguments": "{\"query\":{\"city\":\"NYC\",\"filters\":[\"alerts\",\"daily\"]},\"limit\":3}"
                            }
                        }]
                    },
                    "finish_reason": "tool_calls"
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 5}
            }"#,
            expected: vec![rskit_ai::ToolUseBlock {
                id: "call_1".into(),
                name: "search".into(),
                input: json!({
                    "query": {"city": "NYC", "filters": ["alerts", "daily"]},
                    "limit": 3
                })
                .as_object()
                .cloned()
                .unwrap(),
            }],
        },
    ];

    for case in cases {
        let response = OpenAiDialect::parse_response(case.body).unwrap();
        assert_eq!(response.message.tool_calls, case.expected, "{}", case.name);
    }
}

#[test]
fn parses_streaming_deltas_into_tool_use_blocks() {
    let expected = vec![rskit_ai::ToolUseBlock {
        id: "call_1".into(),
        name: "get_weather".into(),
        input: json!({"location": "NYC", "unit": "c"})
            .as_object()
            .cloned()
            .unwrap(),
    }];

    let actual = streamed_tool_uses(&[
        br#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NY"}}]},"finish_reason":null}]}"#,
        br#"{"choices":[{"delta":{"tool_calls":[{"index":0,"type":"function","function":{"arguments":"C\",\"unit\":\"c\"}"}}]},"finish_reason":null}]}"#,
        br#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#,
    ]);

    assert_eq!(actual, expected);
}