diff --git a/zig/src/server/tool_closers.zig b/zig/src/server/tool_closers.zig new file mode 100644 index 0000000000..dee299c1c8 --- /dev/null +++ b/zig/src/server/tool_closers.zig @@ -0,0 +1,40 @@ +//! Where a reply's namespaced tool-call closers ```` are, found in one pass and kept by NAME. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// A byte a namespaced opener's name may hold after its first. +pub fn nameByte(ch: u8) bool { + return std.ascii.isAlphanumeric(ch) or ch == '_' or ch >= 0x80 or ch == '.' or ch == '-'; +} + +/// One name's closers in order, and how many of them the searches have passed. +const Seen = struct { at: std.ArrayList(usize) = .empty, passed: usize = 0 }; + +pub const Closers = struct { + text: []const u8, + by_name: ?std.StringHashMapUnmanaged(Seen) = null, // read at the first lookup: none without a namespaced opener + + /// Every ```` in `text`, NAME lowered (ASCII only, as ``findIgnoreCasePos`` compares). + fn read(a: Allocator, text: []const u8) Allocator.Error!std.StringHashMapUnmanaged(Seen) { + var by_name: std.StringHashMapUnmanaged(Seen) = .empty; + var pos: usize = 0; + while (std.mem.indexOfPos(u8, text, pos, "")) continue; + const got = try by_name.getOrPut(a, try std.ascii.allocLowerString(a, text[s + 2 .. j])); + if (!got.found_existing) got.value_ptr.* = .{}; + try got.value_ptr.at.append(a, s); + } + return by_name; + } + + /// Where ```` (any case) first starts at or after `from`; `from` must not go down between calls. + pub fn next(c: *Closers, a: Allocator, name: []const u8, from: usize) Allocator.Error!?usize { + if (c.by_name == null) c.by_name = try read(a, c.text); + const seen = c.by_name.?.getPtr(try std.ascii.allocLowerString(a, name)) orelse return null; + while (seen.passed < seen.at.items.len and seen.at.items[seen.passed] < from) seen.passed += 1; + return if (seen.passed < seen.at.items.len) seen.at.items[seen.passed] else null; + } +}; diff --git a/zig/src/server/tool_parse.zig b/zig/src/server/tool_parse.zig index 40ebcbe130..18c308e2d2 100644 --- a/zig/src/server/tool_parse.zig +++ b/zig/src/server/tool_parse.zig @@ -3,6 +3,7 @@ const std = @import("std"); const json = @import("json"); const ids = @import("ids.zig"); const reply_text = @import("reply_text.zig"); +const tool_closers = @import("tool_closers.zig"); const tool_params = @import("tool_params.zig"); const tool_specs = @import("tool_specs.zig"); const tool_streaming = @import("tool_streaming.zig"); @@ -64,17 +65,18 @@ fn envelopes(a: Allocator, text: []const u8) Allocator.Error![]Envelope { pos = e + 12; } pos = 0; + var closers: tool_closers.Closers = .{ .text = text }; while (std.mem.indexOfScalarPos(u8, text, pos, '<')) |s| { pos = s + 1; var j = s + 1; if (j >= text.len or !(std.ascii.isAlphabetic(text[j]) or text[j] == '_')) continue; j += 1; - while (j < text.len and (word(text[j]) or text[j] == '.' or text[j] == '-')) j += 1; + while (j < text.len and tool_closers.nameByte(text[j])) j += 1; if (!std.ascii.startsWithIgnoreCase(text[j..], ":tool_call>")) continue; - const closer = try std.mem.concat(a, u8, &.{ "" }); - const e = findCI(text, closer, j + 11) orelse continue; - try found.append(a, .{ .start = s, .end = e + closer.len, .payload = strip(text[j + 11 .. e]) }); - pos = e + closer.len; + const e = (try closers.next(a, text[s + 1 .. j], j + 11)) orelse continue; + const end = e + j - s + 12; // ```` + try found.append(a, .{ .start = s, .end = end, .payload = strip(text[j + 11 .. e]) }); + pos = end; } pos = 0; while (find(text, "<|tool_call>", pos)) |s| { diff --git a/zig/src/server/tool_parse_test.zig b/zig/src/server/tool_parse_test.zig index d063ac70f8..95233290a9 100644 --- a/zig/src/server/tool_parse_test.zig +++ b/zig/src/server/tool_parse_test.zig @@ -2,6 +2,7 @@ const std = @import("std"); const json = @import("json"); const tool_parse = @import("tool_parse.zig"); +const tool_closers = @import("tool_closers.zig"); const parse = tool_parse.parse; const closeCall = tool_parse.closeCall; const repairGlmKeys = tool_parse.repairGlmKeys; @@ -213,6 +214,52 @@ const nested_calls = [_][]const u8{ "{\"name\":\"lookup\",\"arguments\":{}", }; +test "namespaced openers without closers read in linear time" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array; + // one namespace repeated, then a new namespace per opener + for ([_]bool{ false, true }) |distinct| { + var ns: [2]u64 = undefined; + for ([_]usize{ 1000, 8000 }, &ns) |n, *best| { + var reply: std.ArrayList(u8) = .empty; + for (0..n) |i| try reply.print(a, "", .{if (distinct) i else 0}); + best.* = std.math.maxInt(u64); + for (0..3) |_| { + var run: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer run.deinit(); + const start = std.Io.Clock.awake.now(std.testing.io).nanoseconds; + const r = try parse(run.allocator(), reply.items, tools, null); + best.* = @min(best.*, @as(u64, @intCast(std.Io.Clock.awake.now(std.testing.io).nanoseconds - start))); + try std.testing.expect(r.calls == null); + } + } + // 8x the openers: near 8x the time, not 64x + try std.testing.expect(ns[1] <= 20 * std.time.ns_per_ms or ns[1] <= 24 * ns[0]); + } +} + +test "a closer found by name is the one a search for it finds" { + var prng: std.Random.DefaultPrng = .init(0x7f4a7c15); + const r = prng.random(); + const pieces = [_][]const u8{ "", "", "", "", "", "", "x", "", "" }; + const names = [_][]const u8{ "a", "A", "ab", "b.c", "B.C", "x\u{e9}", "x\u{c9}", "zz" }; + for (0..5_000) |_| { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var text: std.ArrayList(u8) = .empty; + for (0..r.uintLessThan(usize, 16)) |_| try text.appendSlice(a, pieces[r.uintLessThan(usize, pieces.len)]); + var closers: tool_closers.Closers = .{ .text = text.items }; + // every start in order, every name at each, as the parser asks + for (0..text.items.len + 1) |from| for (names) |name| { + const closer = try std.mem.concat(a, u8, &.{ "" }); + try std.testing.expectEqual(std.ascii.findIgnoreCasePos(text.items, from, closer), try closers.next(a, name, from)); + }; + } +} + test "DSML parameters contain complete and unfinished legacy calls without consuming adjacent calls" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit();