summaryrefslogtreecommitdiff
path: root/libpanto/src/conversation.zig
blob: 0e9a4e9574c16d92be8cfbccbab0ea3bcf1bacde (plain)
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
178
179
180
181
182
183
184
185
const std = @import("std");
const Allocator = std.mem.Allocator;

/// A streaming text buffer used by content blocks.
/// Thin alias over ArrayList(u8) — amortized O(1) appends,
/// no O(n²) re-copying.
pub const TextualBlock = std.ArrayList(u8);

/// Create a TextualBlock with initial content (copies the slice).
pub fn textualBlockFromSlice(alloc: Allocator, slice: []const u8) !TextualBlock {
    var buf: TextualBlock = .empty;
    try buf.appendSlice(alloc, slice);
    return buf;
}

pub const ToolUseBlock = struct {
    id: []const u8,
    name: []const u8,
    input: TextualBlock = .empty,

    pub fn deinit(self: *ToolUseBlock, alloc: Allocator) void {
        alloc.free(self.id);
        alloc.free(self.name);
        self.input.deinit(alloc);
    }
};

pub const ToolResultBlock = struct {
    tool_use_id: []const u8,
    content: TextualBlock = .empty,

    pub fn deinit(self: *ToolResultBlock, alloc: Allocator) void {
        alloc.free(self.tool_use_id);
        self.content.deinit(alloc);
    }
};

pub const ContentBlock = union(enum) {
    Text: TextualBlock,
    Thinking: TextualBlock,
    ToolUse: ToolUseBlock,
    ToolResult: ToolResultBlock,

    pub fn deinit(self: *ContentBlock, alloc: Allocator) void {
        switch (self.*) {
            inline else => |*b| b.deinit(alloc),
        }
    }
};

pub const MessageRole = enum {
    system,
    user,
    assistant,
};

pub const Message = struct {
    role: MessageRole,
    content: std.ArrayList(ContentBlock) = .empty,

    pub fn deinit(self: *Message, alloc: Allocator) void {
        for (self.content.items) |*block| {
            block.deinit(alloc);
        }
        self.content.deinit(alloc);
    }
};

pub const Conversation = struct {
    messages: std.ArrayList(Message) = .empty,
    allocator: Allocator,

    pub fn init(allocator: Allocator) Conversation {
        return .{
            .allocator = allocator,
        };
    }

    pub fn addSystemMessage(self: *Conversation, text: []const u8) !void {
        const tb = try textualBlockFromSlice(self.allocator, text);
        var content: std.ArrayList(ContentBlock) = .empty;
        try content.append(self.allocator, .{ .Text = tb });
        try self.messages.append(self.allocator, .{
            .role = .system,
            .content = content,
        });
    }

    pub fn addUserMessage(self: *Conversation, text: []const u8) !void {
        const tb = try textualBlockFromSlice(self.allocator, text);
        var content: std.ArrayList(ContentBlock) = .empty;
        try content.append(self.allocator, .{ .Text = tb });
        try self.messages.append(self.allocator, .{
            .role = .user,
            .content = content,
        });
    }

    /// Append an assistant message. Ownership of the blocks is transferred
    /// to the conversation; the caller must not deinit them after this call.
    pub fn addAssistantMessage(self: *Conversation, blocks: []const ContentBlock) !void {
        var content: std.ArrayList(ContentBlock) = .empty;
        try content.ensureTotalCapacity(self.allocator, blocks.len);
        for (blocks) |block| {
            content.appendAssumeCapacity(block);
        }
        try self.messages.append(self.allocator, .{
            .role = .assistant,
            .content = content,
        });
    }

    pub fn deinit(self: *Conversation) void {
        for (self.messages.items) |*msg| {
            msg.deinit(self.allocator);
        }
        self.messages.deinit(self.allocator);
    }
};

test "Conversation - add messages and verify content" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("You are a helpful assistant.");
    try conv.addUserMessage("Hello!");
    try conv.addAssistantMessage(&.{
        .{ .Text = try textualBlockFromSlice(allocator, "Hi there!") },
    });

    try std.testing.expectEqual(@as(usize, 3), conv.messages.items.len);

    try std.testing.expectEqual(MessageRole.system, conv.messages.items[0].role);
    try std.testing.expectEqualStrings(
        "You are a helpful assistant.",
        conv.messages.items[0].content.items[0].Text.items,
    );

    try std.testing.expectEqual(MessageRole.user, conv.messages.items[1].role);
    try std.testing.expectEqualStrings("Hello!", conv.messages.items[1].content.items[0].Text.items);

    try std.testing.expectEqual(MessageRole.assistant, conv.messages.items[2].role);
    try std.testing.expectEqualStrings("Hi there!", conv.messages.items[2].content.items[0].Text.items);
}

test "TextualBlock - incremental append" {
    const allocator = std.testing.allocator;

    var tb = TextualBlock.empty;
    defer tb.deinit(allocator);

    try tb.appendSlice(allocator, "Hello");
    try tb.appendSlice(allocator, " world");
    try std.testing.expectEqualStrings("Hello world", tb.items);
}

test "Conversation - deinit frees without leaks" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    try conv.addSystemMessage("system");
    try conv.addUserMessage("user message");
    try conv.addAssistantMessage(&.{
        .{ .Text = try textualBlockFromSlice(allocator, "response") },
    });
    conv.deinit();
}

test "ContentBlock - Thinking variant" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addAssistantMessage(&.{
        .{ .Thinking = try textualBlockFromSlice(allocator, "hmm...") },
        .{ .Text = try textualBlockFromSlice(allocator, "answer") },
    });

    try std.testing.expectEqual(@as(usize, 2), conv.messages.items[0].content.items.len);
    try std.testing.expectEqualStrings("hmm...", conv.messages.items[0].content.items[0].Thinking.items);
    try std.testing.expectEqualStrings("answer", conv.messages.items[0].content.items[1].Text.items);
}