diff --git a/build.zig b/build.zig index 96e5c50..6d187e8 100644 --- a/build.zig +++ b/build.zig @@ -45,6 +45,17 @@ pub fn build(b: *std.Build) void { }, }); + const agent = b.addModule("agent", .{ + .root_source_file = b.path("src/agent/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + .{ .name = "provider", .module = provider }, + .{ .name = "testing", .module = testing }, + }, + }); + // This creates a module, which represents a collection of source files alongside // some compilation options, such as optimization mode and linked system libraries. // Zig modules are the preferred way of making Zig code available to consumers. @@ -66,6 +77,7 @@ pub fn build(b: *std.Build) void { .imports = &.{ .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, + .{ .name = "agent", .module = agent }, }, }); @@ -109,6 +121,7 @@ pub fn build(b: *std.Build) void { .{ .name = "coma", .module = mod }, .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, + .{ .name = "agent", .module = agent }, }, }), }); @@ -145,71 +158,16 @@ pub fn build(b: *std.Build) void { run_cmd.addArgs(args); } - // Creates an executable that will run `test` blocks from the provided module. - // Here `mod` needs to define a target, which is why earlier we made sure to - // set the releative field. - const mod_tests = b.addTest(.{ - .root_module = mod, - }); - - // A run step that will run the test executable. - const run_mod_tests = b.addRunArtifact(mod_tests); - - // Creates an executable that will run `test` blocks from the executable's - // root module. Note that test executables only test one module at a time, - // hence why we have to create two separate ones. - const exe_tests = b.addTest(.{ - .root_module = exe.root_module, - }); - - // A run step that will run the second test executable. - const run_exe_tests = b.addRunArtifact(exe_tests); - - // Creates an executable that will run `test` blocks from the provider module. - const provider_tests = b.addTest(.{ - .root_module = provider, - }); - const run_provider_tests = b.addRunArtifact(provider_tests); - - // Creates an executable that will run `test` blocks from the llm module. - const llm_test_module = b.createModule(.{ - .root_source_file = b.path("src/llm/root.tests.zig"), - .target = target, - .optimize = optimize, - .imports = &.{ - .{ .name = "llm", .module = llm }, - .{ .name = "testing", .module = testing }, - }, - }); - const llm_tests = b.addTest(.{ - .root_module = llm_test_module, - }); - const run_llm_tests = b.addRunArtifact(llm_tests); - - // Creates an executable that will run `test` blocks from the testing module. - const testing_tests = b.addTest(.{ - .root_module = testing, - }); - const run_testing_tests = b.addRunArtifact(testing_tests); - - // A top level step for running all tests. dependOn can be called multiple - // times and since the two run steps do not depend on one another, this will - // make the two of them run in parallel. + // Run all tests in the standard build step + const standard_test_suites = createTestSuites(b, target, optimize); const test_step = b.step("test", "Run tests"); - test_step.dependOn(&run_mod_tests.step); - test_step.dependOn(&run_exe_tests.step); - test_step.dependOn(&run_provider_tests.step); - test_step.dependOn(&run_llm_tests.step); - test_step.dependOn(&run_testing_tests.step); + for (standard_test_suites) |suite| { + const run = b.addRunArtifact(suite); + test_step.dependOn(&run.step); + } const coverage_step = b.step("coverage", "Generate coverage reports using kcov"); - const test_suites = [_]*std.Build.Step.Compile{ - mod_tests, - exe_tests, - provider_tests, - llm_tests, - testing_tests, - }; + const coverage_test_suites = createTestSuites(b, target, .ReleaseSafe); const make_dir = b.addSystemCommand(&.{ "mkdir", "-p", "kcov-out" }); const merge_cover = b.addSystemCommand(&.{ @@ -221,8 +179,10 @@ pub fn build(b: *std.Build) void { "kcov-out/suite_2", "kcov-out/suite_3", "kcov-out/suite_4", + "kcov-out/suite_5", }); - for (test_suites, 0..) |test_exe, i| { + + for (coverage_test_suites, 0..) |suite, i| { const out_dir = b.fmt("kcov-out/suite_{d}", .{i}); const run_cover = b.addSystemCommand(&.{ "kcov", @@ -231,13 +191,14 @@ pub fn build(b: *std.Build) void { out_dir, }); run_cover.step.dependOn(&make_dir.step); - run_cover.addArtifactArg(test_exe); + run_cover.addArtifactArg(suite); merge_cover.step.dependOn(&run_cover.step); } - var clean_args = b.allocator.alloc([]const u8, 2 + test_suites.len) catch @panic("OOM"); + + var clean_args = b.allocator.alloc([]const u8, 2 + coverage_test_suites.len) catch @panic("OOM"); clean_args[0] = "rm"; clean_args[1] = "-rf"; - for (test_suites, 0..) |_, i| { + for (coverage_test_suites, 0..) |_, i| { clean_args[2 + i] = b.fmt("kcov-out/suite_{d}", .{i}); } const clean_cover = b.addSystemCommand(clean_args); @@ -256,3 +217,87 @@ pub fn build(b: *std.Build) void { // Lastly, the Zig build system is relatively simple and self-contained, // and reading its source code will allow you to master it. } + +fn createTestSuites( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.OptimizeMode, +) [6]*std.Build.Step.Compile { + const llm = b.createModule(.{ + .root_source_file = b.path("src/llm/root.zig"), + .target = target, + .optimize = optimize, + }); + + const testing = b.createModule(.{ + .root_source_file = b.path("src/testing/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + }, + }); + + const provider = b.createModule(.{ + .root_source_file = b.path("src/provider/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + .{ .name = "testing", .module = testing }, + }, + }); + + const agent = b.createModule(.{ + .root_source_file = b.path("src/agent/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + .{ .name = "provider", .module = provider }, + .{ .name = "testing", .module = testing }, + }, + }); + + const coma = b.createModule(.{ + .root_source_file = b.path("src/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + .{ .name = "provider", .module = provider }, + .{ .name = "agent", .module = agent }, + }, + }); + + const main_module = b.createModule(.{ + .root_source_file = b.path("src/main.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "coma", .module = coma }, + .{ .name = "llm", .module = llm }, + .{ .name = "provider", .module = provider }, + .{ .name = "agent", .module = agent }, + }, + }); + + const llm_test_module = b.createModule(.{ + .root_source_file = b.path("src/llm/root.tests.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "llm", .module = llm }, + .{ .name = "testing", .module = testing }, + }, + }); + + return .{ + b.addTest(.{ .root_module = coma }), + b.addTest(.{ .root_module = main_module }), + b.addTest(.{ .root_module = provider }), + b.addTest(.{ .root_module = agent }), + b.addTest(.{ .root_module = llm_test_module }), + b.addTest(.{ .root_module = testing }), + }; +} diff --git a/src/agent/Agent.zig b/src/agent/Agent.zig new file mode 100644 index 0000000..3db4546 --- /dev/null +++ b/src/agent/Agent.zig @@ -0,0 +1,735 @@ +const std = @import("std"); +const llm = @import("llm"); +const Tool = @import("./Tool.zig"); +const types = @import("./types.zig"); + +const Io = std.Io; +const Allocator = std.mem.Allocator; +const Provider = llm.Provider; +const Future = std.Io.Future; + +const Agent = @This(); + +allocator: Allocator, +io: Io, +provider: Provider, +tools: []const Tool, +session_config: llm.types.SessionConfig, +prev_continuation: ?llm.types.StepContinuation, + +const ToolError = error{ToolNotFound} || Tool.CallError; +pub const AgentError = ToolError || Provider.ProviderError; + +/// Initializes a new Agent instance. +/// +/// `allocator` is used for all internal dynamic memory allocations. +/// `io` is the I/O context to use for operations. +/// `provider` is a Provider interface implementation that determines the LLM provider to use. +/// `tools` is the set of tools available for the agent to execute. +/// `session_config` contains the initial configuration for the LLM session. +pub fn init(allocator: Allocator, io: Io, provider: Provider, tools: []const Tool, session_config: llm.types.SessionConfig) Agent { + return Agent{ + .allocator = allocator, + .io = io, + .provider = provider, + .tools = tools, + .session_config = session_config, + .prev_continuation = null, + }; +} + +/// Deinitializes the Agent, releasing any accumulated session history and internal resources. +pub fn deinit(self: *Agent) void { + if (self.prev_continuation) |*ls| { + ls.deinit(); + self.prev_continuation = null; + } +} + +/// Executes a single non-streaming turn of the agent. +/// +/// The agent sends the turn's prompt to the LLM, handles any tool calls recommended by the model +/// sequentially/concurrently, and returns a `TurnResult` containing the final output and history +/// when the model is finished thinking and using tools. +/// +/// `turn` contains the prompt to send to the LLM. +/// +/// The caller is responsible for deinitializing the returned `TurnResult` by calling deinit() on it. +pub fn executeTurn(self: *Agent, turn: types.Turn) AgentError!types.TurnResult { + return self.executeTurnInternal(turn, null); +} + +/// Executes a single turn of the agent while streaming progress back via a callback. +/// +/// Model chunks and tool results are streamed back via `callback`. +/// Like `executeTurn`, this handles intermediate tool executions, and returns a final `TurnResult`. +/// +/// `turn` contains the prompt to send to the LLM. +/// `callback` is called with the `callback_context` whenever a new streaming chunk or tool execution result is available. +/// `callback_context` is arbitrary data to be passed to the callback as a opaque pointer. +/// +/// The caller is responsible for deinitializing the returned `TurnResult` by calling deinit() on it. +pub fn executeTurnStreaming( + self: *Agent, + turn: types.Turn, + callback: types.StreamingCallback, + callback_context: ?*anyopaque, +) AgentError!types.TurnResult { + var agent_streaming_ctx: StreamingContext = .{ .callback = callback, .context = callback_context }; + return self.executeTurnInternal(turn, &agent_streaming_ctx); +} + +const StreamingContext = struct { + callback: types.StreamingCallback, + context: ?*anyopaque, +}; + +fn streamingCallbackProxy(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) void { + const streaming_ctx: *StreamingContext = @ptrCast(@alignCast(ctx)); + streaming_ctx.callback(streaming_ctx.context, .{ .model_chunk = chunk }); +} + +fn executeToolCall(self: *Agent, tool_call: llm.types.ToolCall) ToolError!llm.types.ToolResult { + const tool = for (self.tools) |t| { + if (std.mem.eql(u8, t.descriptor.name, tool_call.name)) { + break t; + } + } else { + return ToolError.ToolNotFound; + }; + + return try tool.execute(self.allocator, self.io, tool_call.id, tool_call.arguments); +} + +fn executeTurnInternal(self: *Agent, turn: types.Turn, callback_context: ?*StreamingContext) AgentError!types.TurnResult { + var next_steps: std.ArrayList(llm.types.Step) = .empty; + const allocator = self.allocator; + const io = self.io; + defer next_steps.deinit(allocator); + try next_steps.append(allocator, .{ .prompt = turn.prompt }); + + var intermediate_results: std.ArrayList(types.IntermediateStepResult) = .empty; + defer { + // `intermediate_results` will own the memory for all non-final steps; that is, all steps + // that are model or tool outputs but not the initial user input or final result. + // These values will, in the event of a non-error, have their ownership transfered to + // the returned `TurnResult`. + for (intermediate_results.items) |*ir| { + ir.deinit(); + } + intermediate_results.deinit(allocator); + } + + while (true) { + const step_outcome = if (callback_context) |cb| + try self.provider.executeStepStreaming( + allocator, + self.session_config, + next_steps.items, + self.prev_continuation, + streamingCallbackProxy, + cb, + ) + else + try self.provider.executeStep( + allocator, + self.session_config, + next_steps.items, + self.prev_continuation, + ); + next_steps.clearRetainingCapacity(); + + var step_result = step_outcome.result; + const step_continuation = step_outcome.continuation; + + if (self.prev_continuation) |*old_continuation| old_continuation.deinit(); + self.prev_continuation = step_continuation; + + if (step_result.tool_calls.len > 0) { + intermediate_results.append(allocator, .{ .step_result = step_result }) catch |err| { + step_result.deinit(); + return err; + }; + + const SingleToolResult = union(enum) { result: ToolError!llm.types.ToolResult }; + const tool_futures_buf: []SingleToolResult = try allocator.alloc(SingleToolResult, step_result.tool_calls.len); + defer allocator.free(tool_futures_buf); + var tool_futures = Io.Select(SingleToolResult).init(io, tool_futures_buf); + defer while (tool_futures.cancel()) |tool_result| { + var curr_result = tool_result; + if (curr_result.result) |*tr| { + tr.deinit(); + } else |_| {} + }; + for (step_result.tool_calls) |tool_call| { + tool_futures.async(.result, executeToolCall, .{ self, tool_call }); + } + for (0..step_result.tool_calls.len) |_| { + const tool_result_wrapper = tool_futures.await() catch |err| switch (err) { + error.Canceled => unreachable, + }; + var tool_result = try tool_result_wrapper.result; + + intermediate_results.append(allocator, .{ .tool_result = tool_result }) catch |err| { + tool_result.deinit(); + return err; + }; + try next_steps.append(allocator, .{ .tool_result = tool_result }); + + if (callback_context) |cb| { + cb.callback(cb.context, .{ .tool_result = tool_result }); + } + } + } else { + return .{ + .allocator = allocator, + .final_step = step_result, + .intermediate_steps = try intermediate_results.toOwnedSlice(allocator), + }; + } + } +} + +const testing_pkg = @import("testing"); +var test_mock_provider: ?*testing_pkg.MockProvider = null; + +test "Agent.executeTurn - no tool calls" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing_pkg.MockProvider{}; + const prov = mock_provider.provider(); + + const mock_model = llm.types.Model{ + .id = "mock-model", + .display_name = "Mock Model", + }; + + var agent = Agent{ + .allocator = allocator, + .io = io, + .provider = prov, + .tools = &.{}, + .session_config = .{ + .model = mock_model, + .tools = &.{}, + }, + .prev_continuation = null, + }; + defer agent.deinit(); + + mock_provider.execute_step_result = llm.types.StepResult{ + .model_output = &.{.{ .text = "Hello user!" }}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }; + var mock_continuation: testing_pkg.MockProvider.MockStepContinuation = .{}; + mock_provider.execute_step_continuation = mock_continuation.stepContinuation(); + + const turn = types.Turn{ .prompt = "Hi agent" }; + + var result = try agent.executeTurn(turn); + defer result.deinit(); + + try std.testing.expectEqual(@as(usize, 1), mock_provider.execute_step_calls); + try std.testing.expect(agent.prev_continuation != null); + try std.testing.expectEqualStrings("Hello user!", result.final_step.model_output[0].text); +} + +test "Agent.executeTurnStreaming - no tool calls" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing_pkg.MockProvider{}; + const prov = mock_provider.provider(); + + const mock_model = llm.types.Model{ + .id = "mock-model", + .display_name = "Mock Model", + }; + + var agent: Agent = .init( + allocator, + io, + prov, + &.{}, + .{ + .model = mock_model, + .tools = &.{}, + }, + ); + defer agent.deinit(); + + mock_provider.execute_step_result = llm.types.StepResult{ + .model_output = &.{.{ .text = "Hello user!" }}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }; + var mock_continuation: testing_pkg.MockProvider.MockStepContinuation = .{}; + mock_provider.execute_step_continuation = mock_continuation.stepContinuation(); + + const turn = types.Turn{ .prompt = "Hi agent" }; + + const DummyContext = struct { + called: bool = false, + }; + var dummy_ctx = DummyContext{}; + + const callback = struct { + fn cb(ctx: ?*anyopaque, chunk: types.StreamingChunk) void { + _ = chunk; + const c: *DummyContext = @ptrCast(@alignCast(ctx)); + c.called = true; + } + }.cb; + + var result = try agent.executeTurnStreaming(turn, callback, &dummy_ctx); + defer result.deinit(); + + try std.testing.expectEqual(@as(usize, 1), mock_provider.execute_step_streaming_calls); + try std.testing.expectEqual(@as(usize, 0), mock_provider.execute_step_calls); + try std.testing.expect(agent.prev_continuation != null); + try std.testing.expectEqualStrings("Hello user!", result.final_step.model_output[0].text); +} + +const MockToolImpl = struct { + pub fn execute(allocator: std.mem.Allocator, val: i64) ![]const u8 { + if (test_mock_provider) |mp| { + mp.execute_step_result = llm.types.StepResult{ + .model_output = &.{.{ .text = "Final output after tool" }}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = mp, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }; + } + return try std.fmt.allocPrint(allocator, "Tool result for {d}", .{val}); + } +}; + +test "Agent.executeTurn - executes tool call and runs again" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing_pkg.MockProvider{}; + const prov = mock_provider.provider(); + test_mock_provider = &mock_provider; + defer test_mock_provider = null; + + const tool_desc = llm.types.Tool{ + .name = "mock_tool", + .description = "A mock tool for testing", + .parameters = &.{ + .{ + .name = "val", + .description = "integer value", + .type = .integer, + .required = true, + }, + }, + }; + + const tool = Tool.init(tool_desc, MockToolImpl.execute); + const tools = &[_]Tool{tool}; + + const mock_model = llm.types.Model{ + .id = "mock-model", + .display_name = "Mock Model", + }; + + var agent: Agent = .init( + allocator, + io, + prov, + tools, + .{ + .model = mock_model, + .tools = &.{tool.descriptor}, + }, + ); + defer agent.deinit(); + + const args = [_]llm.types.Argument{ + .{ .name = "val", .value = .{ .integer = 42 } }, + }; + const tool_calls = [_]llm.types.ToolCall{ + .{ + .id = "call-id-123", + .name = "mock_tool", + .arguments = @constCast(&args), + }, + }; + + mock_provider.execute_step_result = llm.types.StepResult{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &tool_calls, + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }; + mock_provider.execute_step_continuation = llm.types.StepContinuation{ + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_continuation_vtable, + }; + + const turn = types.Turn{ .prompt = "Hi agent, run mock_tool" }; + + var result = try agent.executeTurn(turn); + defer result.deinit(); + + try std.testing.expectEqual(@as(usize, 2), mock_provider.execute_step_calls); + try std.testing.expect(agent.prev_continuation != null); + try std.testing.expectEqualStrings("Final output after tool", result.final_step.model_output[0].text); +} + +test "Agent.executeTurnStreaming - model chunks streaming" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const CustomStreamingProvider = struct { + fn execute_step_streaming( + ptr: *anyopaque, + alloc: std.mem.Allocator, + session_config: llm.types.SessionConfig, + input: []const llm.types.Step, + previous_step: ?llm.types.StepContinuation, + callback: llm.types.StreamingCallback, + callback_context: ?*anyopaque, + ) llm.Provider.ProviderError!llm.types.StepOutcome { + _ = alloc; + _ = session_config; + _ = input; + _ = previous_step; + + const chunk1 = llm.types.StreamingChunk{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .model_output = .{ .text = "Hello " }, + }, + }, + }, + }, + }; + const chunk2 = llm.types.StreamingChunk{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .model_output = .{ .text = "world!" }, + }, + }, + }, + }, + }; + + callback(callback_context, chunk1); + callback(callback_context, chunk2); + + return llm.types.StepOutcome{ + .result = llm.types.StepResult{ + .model_output = &.{.{ .text = "Hello world!" }}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }, + .continuation = llm.types.StepContinuation{ + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_continuation_vtable, + }, + }; + } + }; + + var custom_vtable = llm.Provider.VTable{ + .list_models = undefined, + .execute_step = undefined, + .execute_step_streaming = CustomStreamingProvider.execute_step_streaming, + .deinit = struct { + fn deinit(ptr: *anyopaque) void { + _ = ptr; + } + }.deinit, + }; + + const prov = llm.Provider{ + .ptr = undefined, + .vtable = &custom_vtable, + }; + + var agent = Agent.init( + allocator, + io, + prov, + &.{}, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = &.{}, + }, + ); + defer agent.deinit(); + + const CallbackState = struct { + const Self = @This(); + chunks: std.ArrayList(types.StreamingChunk) = .empty, + allocator: std.mem.Allocator, + + fn init(alloc: std.mem.Allocator) Self { + return .{ .allocator = alloc }; + } + + fn deinit(self: *Self) void { + self.chunks.deinit(self.allocator); + } + + fn cb(ctx: ?*anyopaque, chunk: types.StreamingChunk) void { + const self: *Self = @ptrCast(@alignCast(ctx.?)); + self.chunks.append(self.allocator, chunk) catch {}; + } + }; + + var cb_state = CallbackState.init(allocator); + defer cb_state.deinit(); + + const turn = types.Turn{ .prompt = "Hi agent" }; + var result = try agent.executeTurnStreaming(turn, CallbackState.cb, &cb_state); + defer result.deinit(); + + try std.testing.expectEqual(@as(usize, 2), cb_state.chunks.items.len); + try std.testing.expectEqualStrings("Hello ", cb_state.chunks.items[0].model_chunk.event.step_event.event.delta.model_output.text); + try std.testing.expectEqualStrings("world!", cb_state.chunks.items[1].model_chunk.event.step_event.event.delta.model_output.text); + try std.testing.expectEqualStrings("Hello world!", result.final_step.model_output[0].text); +} + +test "Agent.executeTurnStreaming - with tool calls" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const CustomStreamingProvider = struct { + const Self = @This(); + calls: usize = 0, + + var args_buf = [_]llm.types.Argument{ + .{ .name = "val", .value = .{ .integer = 42 } }, + }; + var tool_calls_buf = [_]llm.types.ToolCall{ + .{ + .id = "call-id-123", + .name = "mock_tool", + .arguments = &args_buf, + }, + }; + + fn execute_step_streaming( + ptr: *anyopaque, + alloc: std.mem.Allocator, + session_config: llm.types.SessionConfig, + input: []const llm.types.Step, + previous_step: ?llm.types.StepContinuation, + callback: llm.types.StreamingCallback, + callback_context: ?*anyopaque, + ) llm.Provider.ProviderError!llm.types.StepOutcome { + const self: *Self = @ptrCast(@alignCast(ptr)); + self.calls += 1; + _ = callback; + _ = callback_context; + _ = alloc; + _ = session_config; + _ = input; + _ = previous_step; + + if (self.calls == 1) { + return llm.types.StepOutcome{ + .result = llm.types.StepResult{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &tool_calls_buf, + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }, + .continuation = llm.types.StepContinuation{ + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_continuation_vtable, + }, + }; + } else { + return llm.types.StepOutcome{ + .result = llm.types.StepResult{ + .model_output = &.{.{ .text = "Tool executed!" }}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }, + .continuation = llm.types.StepContinuation{ + .ptr = ptr, + .vtable = &testing_pkg.MockProvider.mock_continuation_vtable, + }, + }; + } + } + }; + + var custom_prov_impl = CustomStreamingProvider{}; + var custom_vtable = llm.Provider.VTable{ + .list_models = undefined, + .execute_step = undefined, + .execute_step_streaming = CustomStreamingProvider.execute_step_streaming, + .deinit = struct { + fn deinit(ptr: *anyopaque) void { + _ = ptr; + } + }.deinit, + }; + + const prov = llm.Provider{ + .ptr = &custom_prov_impl, + .vtable = &custom_vtable, + }; + + const tool_desc = llm.types.Tool{ + .name = "mock_tool", + .description = "A mock tool for testing", + .parameters = &.{ + .{ + .name = "val", + .description = "integer value", + .type = .integer, + .required = true, + }, + }, + }; + const tool = Tool.init(tool_desc, MockToolImpl.execute); + const tools = &[_]Tool{tool}; + + var agent = Agent.init( + allocator, + io, + prov, + tools, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = &.{tool_desc}, + }, + ); + defer agent.deinit(); + + const CallbackState = struct { + const Self = @This(); + chunks: std.ArrayList(types.StreamingChunk) = .empty, + allocator: std.mem.Allocator, + + fn init(alloc: std.mem.Allocator) Self { + return .{ .allocator = alloc }; + } + + fn deinit(self: *Self) void { + self.chunks.deinit(self.allocator); + } + + fn cb(ctx: ?*anyopaque, chunk: types.StreamingChunk) void { + const self: *Self = @ptrCast(@alignCast(ctx.?)); + self.chunks.append(self.allocator, chunk) catch {}; + } + }; + + var cb_state = CallbackState.init(allocator); + defer cb_state.deinit(); + + const turn = types.Turn{ .prompt = "Hi agent" }; + var result = try agent.executeTurnStreaming(turn, CallbackState.cb, &cb_state); + defer result.deinit(); + + try std.testing.expectEqual(@as(usize, 1), cb_state.chunks.items.len); + try std.testing.expect(cb_state.chunks.items[0] == .tool_result); + try std.testing.expectEqualStrings("mock_tool", cb_state.chunks.items[0].tool_result.tool_name); + try std.testing.expectEqualStrings("call-id-123", cb_state.chunks.items[0].tool_result.id); + try std.testing.expectEqualStrings("Tool result for 42", cb_state.chunks.items[0].tool_result.result); +} + +test "Agent.executeToolCall - tool not found" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing_pkg.MockProvider{}; + const prov = mock_provider.provider(); + + var agent = Agent.init( + allocator, + io, + prov, + &.{}, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = &.{}, + }, + ); + defer agent.deinit(); + + const tool_call = llm.types.ToolCall{ + .id = "call-id", + .name = "non_existent_tool", + .arguments = &.{}, + }; + + try std.testing.expectError(error.ToolNotFound, agent.executeToolCall(tool_call)); +} + +test "Agent.executeTurn - tool call error cleanup" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing_pkg.MockProvider{}; + const prov = mock_provider.provider(); + + const tool_desc = llm.types.Tool{ + .name = "error_tool", + .description = "A tool that returns error", + .parameters = &.{}, + }; + const error_tool_impl = struct { + fn execute(alloc: Allocator) ![]const u8 { + _ = alloc; + return error.ArgumentTypeMismatch; + } + }; + const tool = Tool.init(tool_desc, error_tool_impl.execute); + const tools = &[_]Tool{tool}; + + var agent = Agent.init( + allocator, + io, + prov, + tools, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = &.{tool_desc}, + }, + ); + defer agent.deinit(); + + const tool_calls = [_]llm.types.ToolCall{ + .{ + .id = "call-id-123", + .name = "error_tool", + .arguments = &.{}, + }, + }; + mock_provider.execute_step_result = llm.types.StepResult{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &tool_calls, + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_step_vtable, + }; + mock_provider.execute_step_continuation = llm.types.StepContinuation{ + .ptr = &mock_provider, + .vtable = &testing_pkg.MockProvider.mock_continuation_vtable, + }; + + const turn = types.Turn{ .prompt = "Run error_tool" }; + try std.testing.expectError(error.ArgumentTypeMismatch, agent.executeTurn(turn)); +} diff --git a/src/agent/Tool.zig b/src/agent/Tool.zig new file mode 100644 index 0000000..6533d81 --- /dev/null +++ b/src/agent/Tool.zig @@ -0,0 +1,878 @@ +const std = @import("std"); +const llm = @import("llm"); + +const Allocator = std.mem.Allocator; +const Io = std.Io; +const Argument = llm.types.Argument; +const ToolResult = llm.types.ToolResult; + +/// Errors that can occur when executing or calling a tool. +pub const CallError = error{ + /// A required argument was not provided. + RequiredArgumentMissing, + /// An argument was provided with a type that does not match the expected type. + ArgumentTypeMismatch, + /// An argument was provided with a name that does not match the expected name; + /// either an unexpected name or the order does not match. + ArgumentMismatch, +} || std.mem.Allocator.Error; + +const Tool = @This(); +const ToolExecuteFn = *const fn (allocator: Allocator, io: Io, args: []const Argument) CallError![]const u8; + +descriptor: llm.types.Tool, +execute_fn: ToolExecuteFn, + +/// Executes the tool with the given arguments. +/// +/// `allocator` is used to allocate memory for the result. +/// `io` is the IO to use for the tool call. +/// `id` is the identifier of the tool call. +/// `args` is the list of arguments to pass to the tool function. All required arguments must be present. +/// Order is not relevant. Unexpected arguments are ignored. +/// +/// Returns a `ToolResult` containing the result of the tool call. The caller is responsible +/// for freeing the `ToolResult` by calling `deinit()`. +pub fn execute(self: *const Tool, allocator: Allocator, io: Io, id: []const u8, args: []const Argument) CallError!ToolResult { + const result = try self.execute_fn(allocator, io, args); + errdefer allocator.free(result); + return ToolResult.initTakingResultOwnership(allocator, self.descriptor.name, id, result); +} + +/// Creates a Tool from a descriptor and a function. +/// Tool calls will delegate to the provided function when called. +/// +/// `descriptor` is the tool's descriptor defining the structure of the tool for the LLM. +/// `execute_fn` is the function to be called when the tool is executed. +/// +/// The function arguments must match the descriptor parameters and be in the same order. +/// The function should also take an allocator as an argument which will be used, at least, +/// to allocate its result. The allocator can be in any position provided the other arguments +/// are still in the same order as the parameters in the descriptor. +/// +/// Additionally, the function can optionally accept an Io struct which represents the IO to use +/// during the tool call. +/// +/// The `execute_fn` should return the result of the tool call as a string and transfer +/// ownership of the memory to the caller. The result will be passed to the LLM as the +/// result of the tool call. +pub fn init(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) Tool { + return comptime blk: { + for (descriptor.parameters, 0..) |p1, i| { + for (descriptor.parameters, 0..) |p2, j| { + if (i != j and std.mem.eql(u8, p1.name, p2.name)) { + @compileError("Duplicate parameter name: " ++ p1.name); + } + } + } + + const res = makeExecuteFn(descriptor, execute_fn); + break :blk Tool{ + .descriptor = descriptor, + .execute_fn = switch (res) { + .ok => |f| f, + .err => |e| @compileError(e.msg), + }, + }; + }; +} + +const ValidationError = error{ + UnsupportedType, + ExpectedFunctionOrPointer, + ArgumentTypeMismatch, + ArgumentCountMismatch, + ReturnTypeMismatch, +}; + +const ValidationResult = union(enum) { + ok: ToolExecuteFn, + err: struct { + code: ValidationError, + msg: []const u8, + }, +}; + +fn expectedTagForType(comptime T: type) ValidationError!struct { tag: std.meta.Tag(Argument.Value), required: bool } { + var InputT = T; + var is_optional = false; + var tag: std.meta.Tag(Argument.Value) = undefined; + + if (@typeInfo(InputT) == .optional) { + is_optional = true; + InputT = @typeInfo(InputT).optional.child; + } + + if (InputT == i64) { + tag = .integer; + } else if (InputT == f64) { + tag = .float; + } else if (InputT == []const u8) { + tag = .string; + } else if (InputT == bool) { + tag = .boolean; + } else { + return ValidationError.UnsupportedType; + } + return .{ .tag = tag, .required = !is_optional }; +} + +fn ExpectedTypeForParam(comptime param: llm.types.Tool.Param) ValidationError!type { + const ExpectedType = switch (param.type) { + .string => []const u8, + .enumeration => []const u8, + .integer => i64, + .float => f64, + .boolean => bool, + .array => return error.UnsupportedType, + }; + + if (!param.required) { + return ?ExpectedType; + } + + return ExpectedType; +} + +/// Looks up an argument by name using a linear search. +/// +/// A linear search is preferred over a hash map lookup here because: +/// 1. The number of arguments passed to a tool call is typically very small (usually < 10). +/// 2. It avoids the overhead and memory allocations of initializing a HashMap. +/// 3. It prevents potential OutOfMemory errors during argument lookup. +inline fn findArgument(args: []const Argument, name: []const u8) ?*const Argument { + for (args) |*arg| { + if (std.mem.eql(u8, arg.name, name)) { + return arg; + } + } + return null; +} + +fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) ValidationResult { + const FnType = @TypeOf(execute_fn); + const fn_info = switch (@typeInfo(FnType)) { + .@"fn" => |info| info, + .pointer => |ptr_info| switch (@typeInfo(ptr_info.child)) { + .@"fn" => |info| info, + else => return .{ .err = .{ .code = ValidationError.ExpectedFunctionOrPointer, .msg = "Expected function or function pointer" } }, + }, + else => return .{ .err = .{ .code = ValidationError.ExpectedFunctionOrPointer, .msg = "Expected function or function pointer" } }, + }; + + const return_type = fn_info.return_type.?; + const payload_type = switch (@typeInfo(return_type)) { + .error_union => |err_union| err_union.payload, + else => return_type, + }; + if (payload_type != []const u8) return .{ .err = .{ .code = ValidationError.ReturnTypeMismatch, .msg = "Function return type must be a string." } }; + + const types = comptime blk: { + var param_idx: usize = 0; + const descriptor_params = descriptor.parameters; + var result_types: [fn_info.params.len]type = undefined; + for (fn_info.params, 0..) |fn_param, i| { + result_types[i] = fn_param.type.?; + + if (fn_param.type.? != Allocator and fn_param.type.? != Io) { + if (param_idx >= descriptor_params.len) { + return .{ .err = .{ .code = ValidationError.ArgumentCountMismatch, .msg = "More arguments in function than descriptor." } }; + } + + const ExpectedType = ExpectedTypeForParam(descriptor_params[param_idx]) catch |val_err| return .{ .err = .{ .code = val_err, .msg = "Failed to resolve type from descriptor (tool: " ++ descriptor.name ++ ", param: " ++ descriptor_params[param_idx].name ++ ")" } }; + if (ExpectedType != fn_param.type.?) { + return .{ .err = .{ + .code = ValidationError.ArgumentTypeMismatch, + .msg = "Argument type mismatch in tool " ++ descriptor.name ++ " for " ++ + "argument " ++ descriptor_params[param_idx].name ++ ": " ++ + "expected " ++ @typeName(ExpectedType) ++ " but got " ++ @typeName(fn_param.type.?), + } }; + } + param_idx += 1; + } + } + + if (param_idx < descriptor_params.len) { + return .{ .err = .{ .code = ValidationError.ArgumentCountMismatch, .msg = "Fewer arguments in function than descriptor." } }; + } + break :blk result_types; + }; + + const TupleType = @Tuple(&types); + return .{ + .ok = struct { + pub fn call(allocator: Allocator, io: Io, input_args: []const Argument) CallError![]const u8 { + var args: TupleType = undefined; + var descriptor_idx: usize = 0; + inline for (0..fn_info.params.len) |func_idx| { + const T = types[func_idx]; + if (T == Allocator) { + args[func_idx] = allocator; + } else if (T == Io) { + args[func_idx] = io; + } else { + const curr_descriptor = descriptor.parameters[descriptor_idx]; + descriptor_idx += 1; + if (findArgument(input_args, curr_descriptor.name)) |argument| { + const expected_tag = comptime try expectedTagForType(T); + if (argument.value != expected_tag.tag) { + return CallError.ArgumentTypeMismatch; + } + args[func_idx] = switch (expected_tag.tag) { + .integer => @intCast(argument.value.integer), + .float => argument.value.float, + .string => argument.value.string, + .boolean => argument.value.boolean, + }; + } else { + if (curr_descriptor.required) return CallError.RequiredArgumentMissing; + if (@typeInfo(@TypeOf(args[func_idx])) != .optional) unreachable; + args[func_idx] = null; + } + } + } + + return @call(.auto, execute_fn, args); + } + }.call, + }; +} + +test init { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "example_function", + .description = "An example function that takes two arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument to the function", + .type = .integer, + .required = true, + }, + .{ + .name = "arg2", + .description = "The second argument to the function", + .type = .string, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn example_function(_: Allocator, arg1: i64, arg2: []const u8) ![]const u8 { + return try std.fmt.allocPrint(allocator, "{d}{s}", .{ arg1, arg2 }); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.example_function); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .integer = 12 } }, + .{ .name = "arg2", .value = .{ .string = "hello" } }, + }; + + var result = try tool.execute(allocator, io, "123", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("example_function", result.tool_name); + try std.testing.expectEqualStrings("123", result.id); + try std.testing.expectEqualStrings("12hello", result.result); +} + +test "makeExecuteFn - ExpectedFunctionOrPointer" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{}, + }; + const res = comptime makeExecuteFn(desc, 42); + try std.testing.expectEqual(ValidationError.ExpectedFunctionOrPointer, res.err.code); +} + +test "makeExecuteFn - ArgumentTypeMismatch" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .integer, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(arg1: []const u8) ![]const u8 { + return arg1; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); + try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected i64 but got []const u8", res.err.msg); +} + +test "makeExecuteFn - ArgumentCountMismatch (too many arguments)" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{}, + }; + const Impl = struct { + pub fn run(arg1: i64) ![]const u8 { + _ = arg1; + return ""; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); + try std.testing.expectEqualStrings("More arguments in function than descriptor.", res.err.msg); +} + +test "makeExecuteFn - ArgumentCountMismatch (too few arguments)" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .integer, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); + try std.testing.expectEqualStrings("Fewer arguments in function than descriptor.", res.err.msg); +} + +test "makeExecuteFn - ParamTypeArrayNotSupported" { + const array_inner_type: llm.types.Tool.Param.Type = .integer; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .{ .array = &array_inner_type }, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(arg1: []const u8) ![]const u8 { + return arg1; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.UnsupportedType, res.err.code); + try std.testing.expectEqualStrings("Failed to resolve type from descriptor (tool: test_tool, param: arg1)", res.err.msg); +} + +test "makeExecuteFn - ReturnTypeMismatch" { + const array_inner_type: llm.types.Tool.Param.Type = .integer; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .{ .array = &array_inner_type }, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(_: []const u8) !i32 { + return 10; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ReturnTypeMismatch, res.err.code); + try std.testing.expectEqualStrings("Function return type must be a string.", res.err.msg); +} + +test "makeExecuteFn - argument optionals OK" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = false, + }, + }, + }; + const Impl = struct { + pub fn run(optional: ?[]const u8) ![]const u8 { + if (optional) |val| return val; + return ""; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expect(res == .ok); +} + +test "makeExecuteFn - argument optional in descriptor, required in fn" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = false, + }, + }, + }; + const Impl = struct { + pub fn run(required: []const u8) ![]const u8 { + return required; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); + try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected ?[]const u8 but got []const u8", res.err.msg); +} + +test "makeExecuteFn - argument required in descriptor, optional in fn" { + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(optional: ?[]const u8) ![]const u8 { + if (optional) |val| return val; + return ""; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); + try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected []const u8 but got ?[]const u8", res.err.msg); +} + +test execute { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: []const u8) ![]const u8 { + return try allocator.dupe(u8, arg1); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{ + .{ + .name = "arg1", + .value = .{ .string = "value" }, + }, + }; + var result = try tool.execute(testing_allocator, io, "id", args); + defer result.deinit(); + + try std.testing.expectEqualStrings("test_tool", result.tool_name); + try std.testing.expectEqualStrings("id", result.id); + try std.testing.expectEqualStrings("value", result.result); +} + +test "execute - unknown argument" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: []const u8) ![]const u8 { + return try allocator.dupe(u8, arg1); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{ + .{ + .name = "unknown", + .value = .{ .string = "value" }, + }, + }; + try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, "id", args)); +} + +test "execute - extra argument ignored" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: []const u8) ![]const u8 { + return try allocator.dupe(u8, arg1); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{ + .{ + .name = "arg1", + .value = .{ .string = "value1" }, + }, + .{ + .name = "arg2", + .value = .{ .string = "value2" }, + }, + }; + var result = try tool.execute(testing_allocator, io, "id", args); + defer result.deinit(); + try std.testing.expectEqualStrings("value1", result.result); +} + +test "execute - missing required argument" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: []const u8) ![]const u8 { + return try allocator.dupe(u8, arg1); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{}; + try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, "id", args)); +} + +test "execute - missing optional argument" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = false, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: ?[]const u8) ![]const u8 { + if (arg1 != null) return allocator.dupe(u8, arg1.?); + return allocator.dupe(u8, "default"); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{}; + var result = try tool.execute(testing_allocator, io, "id", args); + defer result.deinit(); + + try std.testing.expectEqualStrings("test_tool", result.tool_name); + try std.testing.expectEqualStrings("id", result.id); + try std.testing.expectEqualStrings("default", result.result); +} + +test "execute - optional argument provided" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = false, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: ?[]const u8) ![]const u8 { + if (arg1 != null) return allocator.dupe(u8, arg1.?); + return allocator.dupe(u8, "default"); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{ + .{ .name = "arg1", .value = .{ .string = "provided" } }, + }; + var result = try tool.execute(testing_allocator, io, "id", args); + defer result.deinit(); + + try std.testing.expectEqualStrings("test_tool", result.tool_name); + try std.testing.expectEqualStrings("id", result.id); + try std.testing.expectEqualStrings("provided", result.result); +} + +test "execute - argument type mismatch" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + const desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "desc", + .parameters = &.{ + .{ + .name = "arg1", + .description = "desc", + .type = .string, + .required = true, + }, + }, + }; + const Impl = struct { + pub fn run(allocator: Allocator, arg1: []const u8) ![]const u8 { + return try allocator.dupe(u8, arg1); + } + }; + const tool = Tool.init(desc, Impl.run); + const args: []const Argument = &.{ + .{ + .name = "arg1", + .value = .{ .integer = 10 }, + }, + }; + try std.testing.expectError(CallError.ArgumentTypeMismatch, tool.execute(testing_allocator, io, "id", args)); +} + +test "execute - no Allocator parameter" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "no_allocator_func", + .description = "Takes two arguments, no allocator in parameter signature", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .string, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn no_allocator_func(arg1: []const u8) ![]const u8 { + return allocator.dupe(u8, arg1); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.no_allocator_func); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .string = "hello" } }, + }; + + var result = try tool.execute(allocator, io, "abc", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("no_allocator_func", result.tool_name); + try std.testing.expectEqualStrings("abc", result.id); + try std.testing.expectEqualStrings("hello", result.result); +} + +test "execute - Allocator as middle/last parameter" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "middle_last_allocator_func", + .description = "Takes three arguments, allocator is at the middle/end", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .integer, + .required = true, + }, + .{ + .name = "arg2", + .description = "The second argument", + .type = .string, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn middle_last_allocator_func(arg1: i64, tool_allocator: Allocator, arg2: []const u8) ![]const u8 { + return try std.fmt.allocPrint(tool_allocator, "middle-{d}-{s}", .{ arg1, arg2 }); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.middle_last_allocator_func); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .integer = 99 } }, + .{ .name = "arg2", .value = .{ .string = "test" } }, + }; + + var result = try tool.execute(allocator, io, "xyz", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("middle_last_allocator_func", result.tool_name); + try std.testing.expectEqualStrings("xyz", result.id); + try std.testing.expectEqualStrings("middle-99-test", result.result); +} + +test "execute - multiple Allocator parameters" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "multi_allocator_func", + .description = "Takes multiple allocator arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .integer, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn multi_allocator_func(alloc1: Allocator, arg1: i64, alloc2: Allocator) ![]const u8 { + if (alloc1.ptr != alloc2.ptr) return error.OutOfMemory; + return try std.fmt.allocPrint(alloc1, "multi-{d}", .{arg1}); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.multi_allocator_func); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .integer = 7 } }, + }; + + var result = try tool.execute(allocator, io, "multi", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("multi_allocator_func", result.tool_name); + try std.testing.expectEqualStrings("multi", result.id); + try std.testing.expectEqualStrings("multi-7", result.result); +} + +test "execute - Io parameter" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "io_func", + .description = "Takes io and arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .string, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn io_func(tool_io: Io, tool_allocator: Allocator, arg1: []const u8) ![]const u8 { + _ = tool_io; + return try std.fmt.allocPrint(tool_allocator, "io-{s}", .{arg1}); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.io_func); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .string = "test" } }, + }; + + var result = try tool.execute(allocator, io, "io-test", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("io_func", result.tool_name); + try std.testing.expectEqualStrings("io-test", result.id); + try std.testing.expectEqualStrings("io-test", result.result); +} + +test "execute - multiple Io parameters" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + const tool_descriptor: llm.types.Tool = .{ + .name = "multi_io_func", + .description = "Takes multiple io arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .integer, + .required = true, + }, + }, + }; + const tool_impl = struct { + pub fn multi_io_func(io1: Io, arg1: i64, io2: Io, tool_allocator: Allocator) ![]const u8 { + _ = io1; + _ = io2; + return try std.fmt.allocPrint(tool_allocator, "multi-io-{d}", .{arg1}); + } + }; + const tool = comptime init(tool_descriptor, tool_impl.multi_io_func); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .integer = 77 } }, + }; + + var result = try tool.execute(allocator, io, "multi-io", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("multi_io_func", result.tool_name); + try std.testing.expectEqualStrings("multi-io", result.id); + try std.testing.expectEqualStrings("multi-io-77", result.result); +} diff --git a/src/agent/root.zig b/src/agent/root.zig new file mode 100644 index 0000000..72fc1fd --- /dev/null +++ b/src/agent/root.zig @@ -0,0 +1,9 @@ +const std = @import("std"); + +pub const Agent = @import("Agent.zig"); +pub const Tool = @import("Tool.zig"); +pub const types = @import("types.zig"); + +test { + std.testing.refAllDecls(@This()); +} diff --git a/src/agent/types.zig b/src/agent/types.zig new file mode 100644 index 0000000..3efd0f9 --- /dev/null +++ b/src/agent/types.zig @@ -0,0 +1,67 @@ +const std = @import("std"); +const llm = @import("llm"); + +const Allocator = std.mem.Allocator; + +/// Represents the input to start a single execution turn for the Agent. +pub const Turn = struct { + /// The input text prompt or message for this turn. + prompt: []const u8, +}; + +/// The result of an intermediate step executed during an agent turn. +/// +/// An intermediate step is either a response from the model that required further +/// processing (e.g. a tool call) or the result of executing a tool call requested by +/// the model. +pub const IntermediateStepResult = union(enum) { + /// The result of an LLM generation step. + step_result: llm.types.StepResult, + /// The result of executing a tool call requested by the model. + tool_result: llm.types.ToolResult, + + /// Frees the resources associated with the intermediate step result. + pub fn deinit(self: *IntermediateStepResult) void { + switch (self.*) { + .step_result => |*step_result| step_result.deinit(), + .tool_result => |*tool_result| tool_result.deinit(), + } + self.* = undefined; + } +}; + +/// The final result of an agent turn, including all intermediate steps taken and the final LLM step result. +/// Memory must be freed using the `deinit` method. +pub const TurnResult = struct { + /// The allocator used to allocate resources in this TurnResult. + allocator: Allocator, + /// The history of intermediate steps (model thoughts, tool calls, and results) executed during the turn. + intermediate_steps: []IntermediateStepResult, + /// The final step result that completed the turn (typically the model's final text output). + final_step: llm.types.StepResult, + + /// Deinitializes the TurnResult and frees all associated memory. + /// + /// This method must be called exactly once for each TurnResult when it is no longer needed. + /// All `intermediate_steps` and `final_step` will be deinitialized in order to free their + /// allocated resources. + pub fn deinit(self: *TurnResult) void { + for (self.intermediate_steps) |*step| { + step.deinit(); + } + self.final_step.deinit(); + self.allocator.free(self.intermediate_steps); + self.* = undefined; + } +}; + +/// The data chunk representing incremental updates in a streaming Agent turn. +pub const StreamingChunk = union(enum) { + /// An incremental update chunk from the streaming LLM response. + model_chunk: llm.types.StreamingChunk, + /// The outcome/result of executing a tool. + tool_result: llm.types.ToolResult, +}; + +/// Callback function signature for processing streaming response chunks. +pub const StreamingCallback = *const fn (context: ?*anyopaque, chunk: StreamingChunk) void; diff --git a/src/llm/Provider.tests.zig b/src/llm/Provider.tests.zig index 0f0c00f..57c346b 100644 --- a/src/llm/Provider.tests.zig +++ b/src/llm/Provider.tests.zig @@ -31,23 +31,18 @@ test "Provider.executeStep delegates to VTable" { .{ .prompt = "hello" }, }; - const prev_result = llm.types.StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = &mock_impl, - .vtable = &testing.MockProvider.mock_step_vtable, - }; + var mock_step_continuation: testing.MockProvider.MockStepContinuation = .{}; + const last_step_continuation = mock_step_continuation.stepContinuation(); - var result = try prov.executeStep(allocator, session_config, input_steps, prev_result); - defer result.deinit(); + var outcome = try prov.executeStep(allocator, session_config, input_steps, last_step_continuation); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.execute_step_calls); try std.testing.expectEqual(allocator, mock_impl.last_allocator.?); try std.testing.expectEqualStrings("test-model-id", mock_impl.last_session_config.?.model.id); try std.testing.expectEqualStrings("hello", mock_impl.last_input.?[0].prompt); - try std.testing.expectEqual(prev_result.vtable, mock_impl.last_previous_step.?.vtable); - try std.testing.expectEqual(prev_result.ptr, mock_impl.last_previous_step.?.ptr); + try std.testing.expectEqual(last_step_continuation.ptr, mock_impl.last_previous_step.?.ptr); } test "Provider.deinit delegates to VTable" { @@ -109,13 +104,18 @@ test "Provider.executeStep returns custom success and error" { .ptr = &mock_impl, .vtable = &testing.MockProvider.mock_step_vtable, }; + mock_impl.execute_step_continuation = llm.types.StepContinuation{ + .ptr = &mock_impl, + .vtable = &testing.MockProvider.mock_continuation_vtable, + }; var prov = mock_impl.provider(); - var result = try prov.executeStep(allocator, session_config, &.{}, null); - defer result.deinit(); + var outcome = try prov.executeStep(allocator, session_config, &.{}, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.execute_step_calls); - try std.testing.expectEqualStrings("custom-output", result.model_output[0].text); + try std.testing.expectEqualStrings("custom-output", outcome.result.model_output[0].text); } // Test custom error @@ -145,13 +145,8 @@ test "Provider.executeStepStreaming delegates to VTable" { .{ .prompt = "hello" }, }; - const prev_result = llm.types.StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = &mock_impl, - .vtable = &testing.MockProvider.mock_step_vtable, - }; + var mock_continuation: testing.MockProvider.MockStepContinuation = .{}; + const prev_continuation = mock_continuation.stepContinuation(); const CallbackState = struct { fn callback(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) void { @@ -160,16 +155,16 @@ test "Provider.executeStepStreaming delegates to VTable" { } }; - var result = try prov.executeStepStreaming(allocator, session_config, input_steps, prev_result, CallbackState.callback, null); - defer result.deinit(); + var outcome = try prov.executeStepStreaming(allocator, session_config, input_steps, prev_continuation, CallbackState.callback, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.execute_step_streaming_calls); try std.testing.expectEqual(@as(usize, 0), mock_impl.execute_step_calls); try std.testing.expectEqual(allocator, mock_impl.last_allocator.?); try std.testing.expectEqualStrings("test-model-id", mock_impl.last_session_config.?.model.id); try std.testing.expectEqualStrings("hello", mock_impl.last_input.?[0].prompt); - try std.testing.expectEqual(prev_result.vtable, mock_impl.last_previous_step.?.vtable); - try std.testing.expectEqual(prev_result.ptr, mock_impl.last_previous_step.?.ptr); + try std.testing.expectEqual(prev_continuation.ptr, mock_impl.last_previous_step.?.ptr); } test "Provider.executeStepStreaming returns custom success and error" { @@ -197,13 +192,18 @@ test "Provider.executeStepStreaming returns custom success and error" { .ptr = &mock_impl, .vtable = &testing.MockProvider.mock_step_vtable, }; + mock_impl.execute_step_continuation = llm.types.StepContinuation{ + .ptr = &mock_impl, + .vtable = &testing.MockProvider.mock_continuation_vtable, + }; var prov = mock_impl.provider(); - var result = try prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null); - defer result.deinit(); + var outcome = try prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.execute_step_streaming_calls); - try std.testing.expectEqualStrings("custom-output", result.model_output[0].text); + try std.testing.expectEqualStrings("custom-output", outcome.result.model_output[0].text); } // Test custom error @@ -215,4 +215,3 @@ test "Provider.executeStepStreaming returns custom success and error" { try std.testing.expectError(error.HttpRequestFailed, prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null)); } } - diff --git a/src/llm/Provider.zig b/src/llm/Provider.zig index ae42d7e..ddb5749 100644 --- a/src/llm/Provider.zig +++ b/src/llm/Provider.zig @@ -1,11 +1,13 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const types = @import("types.zig"); + const ListModelsResult = types.ListModelsResult; -const StepResult = types.StepResult; const SessionConfig = types.SessionConfig; const Step = types.Step; const StreamingCallback = types.StreamingCallback; +const StepContinuation = types.StepContinuation; +const StepOutcome = types.StepOutcome; /// An interface for an LLM provider. const Provider = @This(); @@ -33,8 +35,8 @@ pub const VTable = struct { allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, - ) ProviderError!StepResult, + previous_step: ?StepContinuation, + ) ProviderError!StepOutcome, /// Executes a single interaction step with the LLM, streaming the response. execute_step_streaming: *const fn ( @@ -42,10 +44,10 @@ pub const VTable = struct { allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, - callback: types.StreamingCallback, + previous_step: ?StepContinuation, + callback: StreamingCallback, callback_context: ?*anyopaque, - ) ProviderError!StepResult, + ) ProviderError!StepOutcome, /// Frees the resources associated with the provider. deinit: *const fn (ptr: *anyopaque) void, @@ -61,25 +63,26 @@ pub fn listModels(provider: *Provider, allocator: Allocator) ProviderError!ListM /// Executes a single step of interaction with the LLM provider. /// -/// `allocator` is used to allocate the structures and strings in the returned `StepResult`. +/// `allocator` is used to allocate the structures and strings in the returned `StepOutcome`. /// `session_config` configuration for the session, should be constant for all calls to executeStep within a session. /// `input` is an array of steps (prompts or tool results) to send to the model. /// `previous_step` is the result of the previous interaction, if any, to maintain context. /// -/// The caller **MUST** call `deinit()` on the returned `StepResult` to free the allocated memory. +/// The caller **MUST** call `deinit()` on both the returned `result` and `continuation` fields of `StepOutcome` +/// to free the allocated memory. pub fn executeStep( provider: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, -) ProviderError!StepResult { + previous_step: ?StepContinuation, +) ProviderError!StepOutcome { return provider.vtable.execute_step(provider.ptr, allocator, session_config, input, previous_step); } /// Executes a single step of interaction with the LLM provider, streaming chunks back to the callback. /// -/// `allocator` is used to allocate internal streaming state and structures in the final returned `StepResult`. +/// `allocator` is used to allocate internal streaming state and structures in the final returned `StepOutcome`. /// `session_config` configuration for the session, should be constant for all calls to executeStep within a session. /// `input` is an array of steps (prompts or tool results) to send to the model. /// `previous_step` is the result of the previous interaction, if any, to maintain context. @@ -89,16 +92,17 @@ pub fn executeStep( /// Memory Behavior: /// - The chunks sent to `callback` are managed by the Provider and will be freed after the callback returns. /// The callback must copy/duplicate any data it needs to retain past the execution of the callback. -/// - The caller **MUST** call `deinit()` on the returned final `StepResult` to free the accumulated response contents. +/// - The caller **MUST** call `deinit()` on both the returned `result` and `continuation` fields of `StepOutcome` +/// to free the accumulated response contents and continuation memory. pub fn executeStepStreaming( provider: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, + previous_step: ?StepContinuation, callback: StreamingCallback, callback_context: ?*anyopaque, -) ProviderError!StepResult { +) ProviderError!StepOutcome { return provider.vtable.execute_step_streaming(provider.ptr, allocator, session_config, input, previous_step, callback, callback_context); } diff --git a/src/llm/root.zig b/src/llm/root.zig index 3e5bbe4..a9ee1c4 100644 --- a/src/llm/root.zig +++ b/src/llm/root.zig @@ -1,2 +1,8 @@ +const std = @import("std"); + pub const Provider = @import("Provider.zig"); pub const types = @import("types.zig"); + +test { + std.testing.refAllDecls(@This()); +} diff --git a/src/llm/types.zig b/src/llm/types.zig index 458a556..15fac0d 100644 --- a/src/llm/types.zig +++ b/src/llm/types.zig @@ -1,5 +1,7 @@ const std = @import("std"); +const Allocator = std.mem.Allocator; + /// Configuration for an LLM interaction session. pub const SessionConfig = struct { /// The model to be used for the session. @@ -103,7 +105,6 @@ pub const ToolCall = struct { /// The result of executing a step (interaction) with the LLM. /// Memory must be freed using the `deinit` method. -/// TODO(razza): Rename to Interaction and rename any related objects/identifiers. pub const StepResult = struct { /// The final output text generated by the model. model_output: []const ModelOutput, @@ -126,6 +127,30 @@ pub const StepResult = struct { } }; +/// Represents information required to continue an interaction with an LLM. +pub const StepContinuation = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + deinit: *const fn (ptr: *anyopaque) void, + }; + + /// Frees the resources associated with the StepContinuation. + pub fn deinit(self: *StepContinuation) void { + self.vtable.deinit(self.ptr); + self.* = undefined; + } +}; + +/// Represents the outcome of executing a single interaction step with the LLM. +pub const StepOutcome = struct { + /// The result of the step execution containing generated outputs, thoughts, and tool calls. + result: StepResult, + /// The continuation required to maintain conversation state for subsequent steps. + continuation: StepContinuation, +}; + /// Represents a single input step to the LLM. pub const Step = union(enum) { /// A user-provided text prompt. @@ -142,6 +167,57 @@ pub const ToolResult = struct { id: []const u8, /// The output or result of the tool execution as a string. result: []const u8, + + allocator: Allocator, + + /// Initializes a new Agent-managed ToolResult and returns a reference to it via `llm.types.ToolResult`. + /// + /// This method duplicates `tool_name`, `id`, and `result` using the provided `allocator` and memory + /// must be released via `deinit()`. + pub fn init(allocator: Allocator, tool_name: []const u8, id: []const u8, result: []const u8) !ToolResult { + const dupe_tool_name = try allocator.dupe(u8, tool_name); + errdefer allocator.free(dupe_tool_name); + const dupe_id = try allocator.dupe(u8, id); + errdefer allocator.free(dupe_id); + const dupe_result = try allocator.dupe(u8, result); + errdefer allocator.free(dupe_result); + + return .{ + .tool_name = dupe_tool_name, + .id = dupe_id, + .result = dupe_result, + .allocator = allocator, + }; + } + + /// Initializes a new Agent-managed ToolResult and returns a reference to it via `llm.types.ToolResult`. + /// + /// This method duplicates `tool_name`, and `id` using the provided allocator but takes ownership of `result`. + /// Memory must be released via `deinit()` and `result` must have been allocated with the same allocator + /// passed to this function. + pub fn initTakingResultOwnership(allocator: Allocator, tool_name: []const u8, id: []const u8, result: []const u8) !ToolResult { + const dupe_tool_name = try allocator.dupe(u8, tool_name); + errdefer allocator.free(dupe_tool_name); + const dupe_id = try allocator.dupe(u8, id); + errdefer allocator.free(dupe_id); + + return .{ + .tool_name = dupe_tool_name, + .id = dupe_id, + .result = result, + .allocator = allocator, + }; + } + + /// Deinitializes the ToolResult and frees all associated memory. + /// + /// This method must be called exactly once for each ToolResult when it is no longer needed. + pub fn deinit(self: *ToolResult) void { + self.allocator.free(self.tool_name); + self.allocator.free(self.id); + self.allocator.free(self.result); + self.* = undefined; + } }; /// The type of a step in a streaming interaction. @@ -166,6 +242,22 @@ pub const StepStartPayload = union(enum) { tool_call: ToolCallStart, }; +/// The incremental argument update payload for a tool call. +pub const ToolCallDelta = struct { + /// The unique ID of the tool call. + /// + /// This is provided for convience and is not part of the delta. It is included to make it + /// easier to track which arguments belong to which tool call. + id: []const u8, + /// The name of the tool. + /// + /// This is provided for convience and is not part of the delta. It is included to make it + /// easier to track which arguments belong to which tool call. + name: []const u8, + /// The arguments delta for tool call parameters, represented as well-typed parsed arguments. + arguments: []const Argument, +}; + /// The incremental update payload. pub const Delta = union(enum) { /// Text content delta from the model output. @@ -173,7 +265,7 @@ pub const Delta = union(enum) { /// Thought content delta from the model's reasoning. thought: Thought, /// Arguments delta for tool call parameters, represented as well-typed parsed arguments. - tool_call: []const Argument, + tool_call: ToolCallDelta, }; /// Payload for a step event in a streaming response. diff --git a/src/main.zig b/src/main.zig index d16b11e..d89ac38 100644 --- a/src/main.zig +++ b/src/main.zig @@ -1,6 +1,10 @@ const std = @import("std"); const provider = @import("provider"); const llm = @import("llm"); +const agent_pkg = @import("agent"); +const Agent = agent_pkg.Agent; +const Tool = agent_pkg.Tool; +const types = agent_pkg.types; const coma = @import("coma"); @@ -114,65 +118,77 @@ fn printMarkdown(stream_ctx: *StreamContext, text: []const u8) void { } } -fn streamCallback(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) void { +fn streamCallback(ctx: ?*anyopaque, agent_chunk: agent_pkg.types.StreamingChunk) void { const stream_ctx: *StreamContext = @ptrCast(@alignCast(ctx)); - switch (chunk.event) { - .interaction_created => {}, - .step_event => |step_ev| { - switch (step_ev.event) { - .start => |start_payload| { - switch (start_payload) { - .thought => { - if (stream_ctx.current_type != .thought) { - std.debug.print("{s}Thinking...{s}\n", .{ color_gray, color_reset }); - stream_ctx.current_type = .thought; - } - }, - .model_output => { - if (stream_ctx.current_type != .model_output) { - std.debug.print("\n{s}Agent >{s} ", .{ color_cyan ++ color_bold, color_reset }); - stream_ctx.current_type = .model_output; - } - }, - .tool_call => |tc| { - if (stream_ctx.current_type != .tool_call) { - std.debug.print("\n{s}[Tool Call: {s}]{s}\n", .{ color_yellow, tc.name, color_reset }); - stream_ctx.current_type = .tool_call; - } - }, - } - }, - .delta => |delta| { - switch (delta) { - .thought => |thought| { - if (stream_ctx.current_type != .thought) { - std.debug.print("{s}Thinking...{s}\n", .{ color_gray, color_reset }); - stream_ctx.current_type = .thought; + switch (agent_chunk) { + .model_chunk => |chunk| { + switch (chunk.event) { + .interaction_created => {}, + .step_event => |step_ev| { + switch (step_ev.event) { + .start => |start_payload| { + switch (start_payload) { + .thought => { + if (stream_ctx.current_type != .thought) { + std.debug.print("{s}Thinking...{s}\n", .{ color_gray, color_reset }); + stream_ctx.current_type = .thought; + } + }, + .model_output => { + if (stream_ctx.current_type != .model_output) { + std.debug.print("\n{s}Agent >{s} ", .{ color_cyan ++ color_bold, color_reset }); + stream_ctx.current_type = .model_output; + } + }, + .tool_call => |tc| { + std.debug.print("\n{s}[Tool Call ({s}): {s}]{s}\n", .{ color_yellow, tc.id, tc.name, color_reset }); + stream_ctx.current_type = .tool_call; + }, } - printMarkdown(stream_ctx, thought.text); }, - .model_output => |mo| { - if (stream_ctx.current_type != .model_output) { - std.debug.print("\n{s}Agent >{s} ", .{ color_cyan ++ color_bold, color_reset }); - stream_ctx.current_type = .model_output; - } - switch (mo) { - .text => |text| { - printMarkdown(stream_ctx, text); + .delta => |delta| { + switch (delta) { + .thought => |thought| { + if (stream_ctx.current_type != .thought) { + std.debug.print("{s}Thinking...{s}\n", .{ color_gray, color_reset }); + stream_ctx.current_type = .thought; + } + printMarkdown(stream_ctx, thought.text); + }, + .model_output => |mo| { + if (stream_ctx.current_type != .model_output) { + std.debug.print("\n{s}Agent >{s} ", .{ color_cyan ++ color_bold, color_reset }); + stream_ctx.current_type = .model_output; + } + switch (mo) { + .text => |text| { + printMarkdown(stream_ctx, text); + }, + } + }, + .tool_call => |dt| { + for (dt.arguments) |arg| { + switch (arg.value) { + .string => |s| std.debug.print(" {s}({s}): \"{s}\"\n", .{ dt.id, arg.name, s }), + .integer => |i| std.debug.print(" {s}({s}): {}\n", .{ dt.id, arg.name, i }), + .float => |f| std.debug.print(" {s}({s}): {d}\n", .{ dt.id, arg.name, f }), + .boolean => |b| std.debug.print(" {s}({s}): {}\n", .{ dt.id, arg.name, b }), + } + } }, } }, - .tool_call => |args| { - _ = args; - }, + .end => {}, } }, - .end => {}, + .interaction_completed => { + flushBackticks(stream_ctx); + flushAsterisks(stream_ctx); + }, } }, - .interaction_completed => { - flushBackticks(stream_ctx); - flushAsterisks(stream_ctx); + .tool_result => |tr| { + std.debug.print("{s}Output ({s}):{s}\n{s}\n", .{ color_green, tr.id, color_reset, tr.result }); }, } } @@ -213,6 +229,38 @@ fn loadApiKey(allocator: std.mem.Allocator, io: std.Io, environ_map: *std.proces return error.ApiKeyMissing; } +fn executeTypescript(allocator: std.mem.Allocator, io: std.Io, code: []const u8) ![]const u8 { + const argv = [_][]const u8{ + "npx", "tsx", "-e", code, + }; + const result = std.process.run(allocator, io, .{ + .argv = &argv, + }) catch |err| { + return try std.fmt.allocPrint(allocator, "Error executing script: {}", .{err}); + }; + errdefer { + allocator.free(result.stdout); + allocator.free(result.stderr); + } + allocator.free(result.stderr); + + if (result.term != .exited) { + allocator.free(result.stdout); + return try allocator.dupe(u8, "Error: process did not exit cleanly"); + } + + return result.stdout; +} + +fn getWeather(allocator: std.mem.Allocator, zip_code: i64) ![]const u8 { + const result_str = if (zip_code == 7302) + try allocator.dupe(u8, "Weather report for 07302: Sunny, 72°F, Humidity 50%, Wind 5 mph") + else + try std.fmt.allocPrint(allocator, "Error: Weather data is only available for zip code 07302. Requested: {}", .{zip_code}); + + return result_str; +} + /// The main entry point of the application. /// Currently used for testing. pub fn main(init: std.process.Init) !void { @@ -245,36 +293,44 @@ pub fn main(init: std.process.Init) !void { } } else unreachable; - const session_config: llm.types.SessionConfig = .{ - .model = selected_model.?, - .tools = &.{ - .{ - .name = "execute_typescript", - .description = "Executes typescript code and returns the output printed to stdout. Takes a single string argument.", - .parameters = &.{ - .{ - .name = "code", - .type = .string, - .required = true, - .description = "The typescript code to execute.", - }, + const tools = &[_]Tool{ + Tool.init(.{ + .name = "execute_typescript", + .description = "Executes typescript code and returns the output printed to stdout. Takes a single string argument.", + .parameters = &.{ + .{ + .name = "code", + .type = .string, + .required = true, + .description = "The typescript code to execute.", }, }, - .{ - .name = "get_weather", - .description = "Get the current weather for a given zip code.", - .parameters = &.{ - .{ - .name = "zip_code", - .type = .integer, - .required = true, - .description = "The 5-digit zip code to get the weather for.", - }, + }, executeTypescript), + Tool.init(.{ + .name = "get_weather", + .description = "Get the current weather for a given zip code.", + .parameters = &.{ + .{ + .name = "zip_code", + .type = .integer, + .required = true, + .description = "The 5-digit zip code to get the weather for.", }, }, + }, getWeather), + }; + + const session_config: llm.types.SessionConfig = .{ + .model = selected_model.?, + .tools = &.{ + tools[0].descriptor, + tools[1].descriptor, }, }; + var agent: Agent = .init(allocator, io, client, tools, session_config); + defer agent.deinit(); + std.debug.print( \\{s}============================================================================ \\ COMA Agent Chat Interface @@ -289,90 +345,22 @@ pub fn main(init: std.process.Init) !void { var stdin_buffer: [1024]u8 = undefined; var stdin_reader = std.Io.File.stdin().reader(io, &stdin_buffer); - // 2. Read until the user hits Enter ('\n') - // This returns a slice of the buffer containing the raw string - var last_step: ?llm.types.StepResult = null; - var last_tool_results: std.ArrayList(llm.types.ToolResult) = .empty; - defer last_tool_results.deinit(allocator); while (true) { - var new_steps: std.ArrayList(llm.types.Step) = .empty; - defer new_steps.deinit(allocator); - - if (last_step == null or last_step.?.tool_calls.len == 0) { - std.debug.print("\n{s}User > {s}", .{ color_green ++ color_bold, color_reset }); - const user_input = try stdin_reader.interface.takeDelimiter('\n') orelse break; - if (user_input.len == 0) break; - try new_steps.append(allocator, .{ .prompt = user_input }); - } else { - for (last_step.?.tool_calls) |tool_call| { - if (std.mem.eql(u8, tool_call.name, "execute_typescript")) { - const code = tool_call.arguments[0].value.string; - std.debug.print("\n{s}Executing TypeScript code...{s}\n", .{ color_yellow, color_reset }); - std.debug.print("{s}--- CODE ---{s}\n{s}\n{s}------------{s}\n", .{ color_gray, color_reset, code, color_gray, color_reset }); - const argv = [_][]const u8{ - "npx", "tsx", "-e", code, - }; - const result = try std.process.run(allocator, io, .{ - .argv = &argv, - }); - - if (result.term == .exited) { - const tool_result: llm.types.ToolResult = .{ - .tool_name = tool_call.name, - .id = tool_call.id, - .result = result.stdout, - }; - std.debug.print("{s}Output:{s}\n{s}", .{ color_green, color_reset, result.stdout }); - try last_tool_results.append(allocator, tool_result); - try new_steps.append(allocator, .{ .tool_result = tool_result }); - allocator.free(result.stderr); - } - } else if (std.mem.eql(u8, tool_call.name, "get_weather")) { - const zip_val = tool_call.arguments[0].value; - const zip = switch (zip_val) { - .integer => |v| v, - else => -1, - }; - const result_str = if (zip == 7302) - try allocator.dupe(u8, "Weather report for 07302: Sunny, 72°F, Humidity 50%, Wind 5 mph") - else - try std.fmt.allocPrint(allocator, "Error: Weather data is only available for zip code 07302. Requested: {}", .{zip}); - - std.debug.print("\n{s}Executing get_weather for zip code {}...{s}\n", .{ color_yellow, zip, color_reset }); - - const tool_result: llm.types.ToolResult = .{ - .tool_name = tool_call.name, - .id = tool_call.id, - .result = result_str, - }; - std.debug.print("{s}Output:{s}\n{s}\n", .{ color_green, color_reset, result_str }); - try last_tool_results.append(allocator, tool_result); - try new_steps.append(allocator, .{ .tool_result = tool_result }); - } - } - } + std.debug.print("\n{s}User > {s}", .{ color_green ++ color_bold, color_reset }); + const user_input = try stdin_reader.interface.takeDelimiter('\n') orelse break; + if (user_input.len == 0) break; + const turn = types.Turn{ .prompt = user_input }; var stream_ctx = StreamContext{ .allocator = allocator }; - const step_result = try client.executeStepStreaming(allocator, session_config, new_steps.items, last_step, streamCallback, &stream_ctx); - errdefer step_result.deinit(); - - for (last_tool_results.items) |tool_call| { - allocator.free(tool_call.result); - } - last_tool_results.clearRetainingCapacity(); + var result = agent.executeTurnStreaming(turn, streamCallback, &stream_ctx) catch |err| { + std.debug.print("Error during execution: {}\n", .{err}); + continue; + }; + defer result.deinit(); if (stream_ctx.current_type != null) { std.debug.print("\n", .{}); } - - if (last_step) |*last| { - last.deinit(); - } - last_step = step_result; - } - - if (last_step) |*last| { - last.deinit(); } } diff --git a/src/provider/google/accumulator.zig b/src/provider/google/accumulator.zig new file mode 100644 index 0000000..d7c385e --- /dev/null +++ b/src/provider/google/accumulator.zig @@ -0,0 +1,455 @@ +const std = @import("std"); +const llm = @import("llm"); +const api = @import("api.zig"); +const converter = @import("converter.zig"); + +const Allocator = std.mem.Allocator; +const ProviderError = llm.Provider.ProviderError; + +/// Accumulator union used to buffer and aggregate streaming responses for each step. +/// Handles text/thought aggregation and tool call parsing during streaming. +pub const StepAccumulator = union(enum) { + /// Holds the aggregated thinking/thought process content. + thought: std.Io.Writer.Allocating, + /// Holds the aggregated model response output text. + model_output: std.Io.Writer.Allocating, + /// Holds tool call details and buffers the arguments JSON stream. + tool_call: struct { + /// The unique ID of the tool call. + id: []const u8, + /// The name of the function to invoke. + name: []const u8, + /// Dynamic writer buffer accumulating arguments JSON text. + arguments_json: std.Io.Writer.Allocating, + /// Tracks the count of arguments that have already been streamed out to the callback. + processed_argument_count: usize, + }, + /// Represents an uninitialized or empty accumulator slot. + empty, + + /// Initializes the `StepAccumulator` with an initial step context. + /// + /// `result_arena_allocator` is used to allocate the internal writers and copy initial values. + /// Memory allocated in `result_arena_allocator` is owned by the caller's session arena. + /// `initial_step` provides the starting type and content (e.g. initial thought, model output, or function call). + /// + /// Returns the initialized `StepAccumulator`, or an allocation error. + pub fn init(result_arena_allocator: Allocator, initial_step: api.Step) !@This() { + switch (initial_step) { + .thought => |contents| { + var list = std.Io.Writer.Allocating.init(result_arena_allocator); + for (contents) |c| { + if (c.text) |t| { + try list.writer.writeAll(t); + } + } + return .{ .thought = list }; + }, + .model_output => |contents| { + var list = std.Io.Writer.Allocating.init(result_arena_allocator); + for (contents) |c| { + if (c.text) |t| { + try list.writer.writeAll(t); + } + } + return .{ .model_output = list }; + }, + .function_call => |call| return .{ + .tool_call = .{ + .id = try result_arena_allocator.dupe(u8, call.id), + .name = try result_arena_allocator.dupe(u8, call.name), + .arguments_json = std.Io.Writer.Allocating.init(result_arena_allocator), + .processed_argument_count = 0, + }, + }, + } + } + + /// Appends stream delta data to the accumulator. + /// + /// `step_self` points to the accumulator instance. + /// `delta` is the stream delta payload containing new text or argument fragments. + /// + /// Returns a `ProviderError.BadResponse` if the delta type does not match the active accumulator variant. + pub fn appendStep(step_self: *@This(), delta: api.InteractionStepDelta) !void { + switch (step_self.*) { + .thought => |*list| { + if (delta != .thought_summary) return ProviderError.BadResponse; + if (delta.thought_summary.content.text) |t| { + try list.writer.writeAll(t); + } + }, + .model_output => |*list| { + if (delta != .text_delta) return ProviderError.BadResponse; + if (delta.text_delta.text) |t| { + try list.writer.writeAll(t); + } + }, + .tool_call => |*tc| { + if (delta != .arguments_delta) return ProviderError.BadResponse; + try tc.arguments_json.writer.writeAll(delta.arguments_delta.arguments); + }, + .empty => { + return ProviderError.BadResponse; + }, + } + } + + /// Processes newly accumulated tool call arguments JSON and parses new complete arguments. + /// + /// `acc` is the active `StepAccumulator` (must be the `.tool_call` variant). + /// `allocator` is used to allocate the returned array of arguments. + /// + /// **Memory Alert**: + /// The caller takes ownership of the returned slice of `llm.types.Argument` and **MUST** free + /// the slice, the `name` strings, and the `value` strings when done. + /// + /// Returns a slice of newly completed `llm.types.Argument` parameters, or `null` if no new arguments + /// were completed or parsing failed. + pub fn handleToolCallDelta( + acc: *StepAccumulator, + allocator: Allocator, + ) !?[]const llm.types.Argument { + const json_parsed_args = std.json.parseFromSlice( + std.json.Value, + allocator, + acc.tool_call.arguments_json.written(), + .{}, + ) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return null; + }; + defer json_parsed_args.deinit(); + + if (json_parsed_args.value != .object) { + return null; + } + + const function_arguments = api.FunctionArgument.parseFromJsonObject( + allocator, + json_parsed_args.value, + ) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return null; + }; + defer allocator.free(function_arguments); + + const new_count = function_arguments.len - acc.tool_call.processed_argument_count; + if (new_count == 0) { + return null; + } + + const arguments = try converter.dupeArguments(allocator, function_arguments[acc.tool_call.processed_argument_count..]); + errdefer { + for (arguments) |arg| { + allocator.free(arg.name); + switch (arg.value) { + .string => |s| allocator.free(s), + else => {}, + } + } + allocator.free(arguments); + } + + acc.tool_call.processed_argument_count = function_arguments.len; + return arguments; + } + + /// Builds a transient delta payload if new contents or completed arguments are available. + /// + /// `self` is the accumulator instance. + /// `allocator` is used to allocate new arguments for a tool call. + /// `delta` is the incoming interaction step delta. + /// + /// Returns a `TransientDelta` containing the mapped LLM delta if new data was produced, + /// or `null` if no new delta is ready to be sent (e.g. no new complete tool call arguments). + pub fn buildDelta( + self: *@This(), + allocator: Allocator, + delta: api.InteractionStepDelta, + ) !?TransientDelta { + return switch (self.*) { + .thought => .{ + .delta = converter.toThoughtDelta(delta.thought_summary), + .allocator = allocator, + }, + .model_output => .{ + .delta = converter.toModelOutputDelta(delta.text_delta), + .allocator = allocator, + }, + .tool_call => |*tc| { + if (try self.handleToolCallDelta(allocator)) |args| { + return .{ + .delta = converter.toToolCallDelta(tc.id, tc.name, args), + .allocator = allocator, + }; + } + return null; + }, + .empty => null, + }; + } +}; + +/// A wrapper around `llm.types.Delta` that manages the lifetime of allocated delta fields. +/// Holds the allocated fields and the allocator used to create them, so they can be +/// cleanly freed on destruction via `deinit`. +pub const TransientDelta = struct { + /// The wrapped LLM delta payload. + delta: llm.types.Delta, + /// The allocator used to allocate fields in `delta` (such as tool call arguments). + allocator: Allocator, + + /// Frees the resources associated with the nested `delta` payload using the stored allocator. + pub fn deinit(self: @This()) void { + switch (self.delta) { + .tool_call => |tc| { + for (tc.arguments) |arg| { + self.allocator.free(arg.name); + switch (arg.value) { + .string => |s| self.allocator.free(s), + else => {}, + } + } + self.allocator.free(tc.arguments); + }, + else => {}, + } + } +}; + +test "StepAccumulator mismatch delta and init thought" { + const allocator = std.testing.allocator; + + var thoughts = [_]api.Content{.{ .type = .text, .text = "thought_1" }}; + const init_step = api.Step{ .thought = &thoughts }; + var acc = try StepAccumulator.init(allocator, init_step); + defer acc.thought.deinit(); + try std.testing.expect(acc == .thought); + try std.testing.expectEqualStrings("thought_1", acc.thought.written()); + + // Append mismatched delta (should fail) + const mismatch_delta = api.InteractionStepDelta{ + .text_delta = .{ .type = .text, .text = "text" }, + }; + try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); +} + +test "StepAccumulator mismatch delta and init model_output" { + const allocator = std.testing.allocator; + + var outputs = [_]api.Content{.{ .type = .text, .text = "output_1" }}; + const init_step = api.Step{ .model_output = &outputs }; + var acc = try StepAccumulator.init(allocator, init_step); + defer acc.model_output.deinit(); + try std.testing.expect(acc == .model_output); + try std.testing.expectEqualStrings("output_1", acc.model_output.written()); + + // Append mismatched delta (should fail) + const mismatch_delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "args" }, + }; + try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); +} + +test "StepAccumulator mismatch delta and init tool_call" { + const allocator = std.testing.allocator; + + const init_step = api.Step{ + .function_call = .{ + .id = "fc_id", + .name = "fc_name", + .arguments = &.{}, + }, + }; + var acc = try StepAccumulator.init(allocator, init_step); + defer { + allocator.free(acc.tool_call.id); + allocator.free(acc.tool_call.name); + acc.tool_call.arguments_json.deinit(); + } + try std.testing.expect(acc == .tool_call); + try std.testing.expectEqualStrings("fc_id", acc.tool_call.id); + + // Append mismatched delta (should fail) + const mismatch_delta = api.InteractionStepDelta{ + .thought_summary = .{ .content = .{ .type = .text, .text = "think" } }, + }; + try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); + + // Test handleToolCallDelta when json is not valid (returns null) + try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); + + // Test append valid delta and handleToolCallDelta returning null when new_count == 0 + const args_delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "{\"arg1\": \"val1\"}" }, + }; + try acc.appendStep(args_delta); + const args = try acc.handleToolCallDelta(allocator); + try std.testing.expect(args != null); + defer { + for (args.?) |arg| { + allocator.free(arg.name); + switch (arg.value) { + .string => |s| allocator.free(s), + else => {}, + } + } + allocator.free(args.?); + } + + // Calling handleToolCallDelta again without new changes should return null + try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); +} + +test "StepAccumulator mismatch delta and init empty" { + var acc: StepAccumulator = .empty; + const delta = api.InteractionStepDelta{ + .text_delta = .{ .type = .text, .text = "text" }, + }; + try std.testing.expectError(error.BadResponse, acc.appendStep(delta)); +} + +test "StepAccumulator buildDelta thought" { + const allocator = std.testing.allocator; + + var thoughts = [_]api.Content{.{ .type = .text, .text = "thought_init" }}; + var acc = try StepAccumulator.init(allocator, .{ .thought = &thoughts }); + defer acc.thought.deinit(); + + const delta = api.InteractionStepDelta{ + .thought_summary = .{ + .content = .{ .type = .text, .text = "thought_more" }, + }, + }; + try acc.appendStep(delta); + const t_delta = (try acc.buildDelta(allocator, delta)).?; + defer t_delta.deinit(); + + try std.testing.expect(t_delta.delta == .thought); + try std.testing.expectEqualStrings("thought_more", t_delta.delta.thought.text); +} + +test "StepAccumulator buildDelta model_output" { + const allocator = std.testing.allocator; + + var outputs = [_]api.Content{.{ .type = .text, .text = "output_init" }}; + var acc = try StepAccumulator.init(allocator, .{ .model_output = &outputs }); + defer acc.model_output.deinit(); + + const delta = api.InteractionStepDelta{ + .text_delta = .{ .type = .text, .text = "output_more" }, + }; + try acc.appendStep(delta); + const t_delta = (try acc.buildDelta(allocator, delta)).?; + defer t_delta.deinit(); + + try std.testing.expect(t_delta.delta == .model_output); + try std.testing.expectEqualStrings("output_more", t_delta.delta.model_output.text); +} + +test "StepAccumulator buildDelta tool_call" { + const allocator = std.testing.allocator; + + var acc = try StepAccumulator.init(allocator, .{ + .function_call = .{ + .id = "tc_boston", + .name = "get_weather", + .arguments = &.{}, + }, + }); + defer { + allocator.free(acc.tool_call.id); + allocator.free(acc.tool_call.name); + acc.tool_call.arguments_json.deinit(); + } + + const delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "{\"location\": \"Boston\"}" }, + }; + try acc.appendStep(delta); + + const t_delta_opt = try acc.buildDelta(allocator, delta); + try std.testing.expect(t_delta_opt != null); + const t_delta = t_delta_opt.?; + defer t_delta.deinit(); + + try std.testing.expect(t_delta.delta == .tool_call); + try std.testing.expectEqualStrings("tc_boston", t_delta.delta.tool_call.id); + try std.testing.expectEqualStrings("get_weather", t_delta.delta.tool_call.name); + try std.testing.expectEqual(1, t_delta.delta.tool_call.arguments.len); + try std.testing.expectEqualStrings("location", t_delta.delta.tool_call.arguments[0].name); + try std.testing.expectEqualStrings("Boston", t_delta.delta.tool_call.arguments[0].value.string); +} + +test "StepAccumulator handleToolCallDelta invalid JSON structures" { + const allocator = std.testing.allocator; + + var acc = try StepAccumulator.init(allocator, .{ + .function_call = .{ + .id = "tc_invalid", + .name = "some_func", + .arguments = &.{}, + }, + }); + defer { + allocator.free(acc.tool_call.id); + allocator.free(acc.tool_call.name); + acc.tool_call.arguments_json.deinit(); + } + + // 1. JSON is not an object (e.g. an array, which is invalid for tool call arguments) + const array_delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "[1, 2, 3]" }, + }; + try acc.appendStep(array_delta); + try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); + + // Clear and reset state to test another scenario + acc.tool_call.arguments_json.clearRetainingCapacity(); + acc.tool_call.processed_argument_count = 0; + + // 2. JSON is an object but contains unsupported nested types (e.g. nested object) + const nested_delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "{\"nested\": {}}" }, + }; + try acc.appendStep(nested_delta); + try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); + + // Clear and reset state to test null/unsupported arrays + acc.tool_call.arguments_json.clearRetainingCapacity(); + acc.tool_call.processed_argument_count = 0; + + const nested_arr_delta = api.InteractionStepDelta{ + .arguments_delta = .{ .arguments = "{\"nested_arr\": [1]}" }, + }; + try acc.appendStep(nested_arr_delta); + try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); +} + +test "TransientDelta deinit all argument value types" { + const allocator = std.testing.allocator; + + // Set up standard arguments of string, integer, float, boolean + var args = try allocator.alloc(llm.types.Argument, 4); + errdefer allocator.free(args); + + args[0] = .{ .name = try allocator.dupe(u8, "arg_str"), .value = .{ .string = try allocator.dupe(u8, "val_str") } }; + args[1] = .{ .name = try allocator.dupe(u8, "arg_int"), .value = .{ .integer = 42 } }; + args[2] = .{ .name = try allocator.dupe(u8, "arg_flt"), .value = .{ .float = 3.14 } }; + args[3] = .{ .name = try allocator.dupe(u8, "arg_bool"), .value = .{ .boolean = true } }; + + const t_delta = TransientDelta{ + .delta = .{ + .tool_call = .{ + .id = "tc_all_types", + .name = "func_all_types", + .arguments = args, + }, + }, + .allocator = allocator, + }; + + // This should cleanly deallocate the string values, names, and the arguments slice itself. + t_delta.deinit(); +} diff --git a/src/provider/google/converter.zig b/src/provider/google/converter.zig index f986e80..db2c650 100644 --- a/src/provider/google/converter.zig +++ b/src/provider/google/converter.zig @@ -155,25 +155,31 @@ pub fn toStepStartPayload(step: api.Step) llm.types.StepStartPayload { }; } -/// Converts a Gemini API step delta into a generic delta. -/// For tool calls, the parsed arguments must be provided. -/// -/// Returns null if no interesting or reasonable `Delta` object can be constructed. For example, -/// if a tool call delta is received but no arguments have been provided yet. -pub fn toDelta(delta: api.InteractionStepDelta, arguments: []const llm.types.Argument) ?llm.types.Delta { - return switch (delta) { - .arguments_delta => if (arguments.len > 0) .{ - .tool_call = arguments, - } else null, - .text_delta => |td| .{ - .model_output = .{ - .text = td.text orelse "", - }, +/// Converts a Gemini API text delta into a generic model output delta. +pub fn toModelOutputDelta(td: api.TextDelta) llm.types.Delta { + return .{ + .model_output = .{ + .text = td.text orelse "", }, - .thought_summary => |ts| .{ - .thought = .{ - .text = ts.content.text orelse "", - }, + }; +} + +/// Converts a Gemini API thought summary into a generic thought delta. +pub fn toThoughtDelta(ts: api.ThoughtSummaryDelta) llm.types.Delta { + return .{ + .thought = .{ + .text = ts.content.text orelse "", + }, + }; +} + +/// Converts generic tool arguments into a generic tool call delta. +pub fn toToolCallDelta(tool_id: []const u8, tool_name: []const u8, arguments: []const llm.types.Argument) llm.types.Delta { + return .{ + .tool_call = .{ + .id = tool_id, + .name = tool_name, + .arguments = arguments, }, }; } @@ -465,12 +471,11 @@ test toGoogleStep { try std.testing.expectEqualStrings("Hello, world!", google_prompt.user_input.content[0].text.?); try std.testing.expect(google_prompt.user_input.content[0].type == api.Content.Type.text); + var tool_result = try llm.types.ToolResult.init(allocator, "test_tool", "call_123", "success"); + defer tool_result.deinit(); + const tool_step = llm.types.Step{ - .tool_result = .{ - .tool_name = "test_tool", - .id = "call_123", - .result = "success", - }, + .tool_result = tool_result, }; const google_tool = try toGoogleStep(arena_allocator, tool_step); try std.testing.expectEqualStrings("test_tool", google_tool.function_result.name); @@ -499,32 +504,31 @@ test toStepStartPayload { try std.testing.expectEqualStrings("func_name", fc_payload.tool_call.name); } -test toDelta { - const arguments = &[_]llm.types.Argument{ - .{ .name = "arg1", .value = .{ .string = "val1" } }, - }; - const arg_delta = api.InteractionStepDelta{ - .arguments_delta = .{ .arguments = "foo" }, - }; - const delta1 = toDelta(arg_delta, arguments); - try std.testing.expect(delta1.? == .tool_call); - try std.testing.expectEqualStrings("arg1", delta1.?.tool_call[0].name); +test toModelOutputDelta { + const text_delta = api.TextDelta{ .type = .text, .text = "hello" }; + const delta = toModelOutputDelta(text_delta); + try std.testing.expect(delta == .model_output); + try std.testing.expectEqualStrings("hello", delta.model_output.text); +} - const text_delta = api.InteractionStepDelta{ - .text_delta = .{ .type = .text, .text = "hello" }, +test toThoughtDelta { + const thought_delta = api.ThoughtSummaryDelta{ + .content = .{ .type = .text, .text = "thinking" }, }; - const delta2 = toDelta(text_delta, &.{}); - try std.testing.expect(delta2.? == .model_output); - try std.testing.expectEqualStrings("hello", delta2.?.model_output.text); + const delta = toThoughtDelta(thought_delta); + try std.testing.expect(delta == .thought); + try std.testing.expectEqualStrings("thinking", delta.thought.text); +} - const thought_delta = api.InteractionStepDelta{ - .thought_summary = .{ - .content = .{ .type = .text, .text = "thinking" }, - }, +test toToolCallDelta { + const arguments = &[_]llm.types.Argument{ + .{ .name = "arg1", .value = .{ .string = "val1" } }, }; - const delta3 = toDelta(thought_delta, &.{}); - try std.testing.expect(delta3.? == .thought); - try std.testing.expectEqualStrings("thinking", delta3.?.thought.text); + const delta = toToolCallDelta("call_id", "func_name", arguments); + try std.testing.expect(delta == .tool_call); + try std.testing.expectEqualStrings("call_id", delta.tool_call.id); + try std.testing.expectEqualStrings("func_name", delta.tool_call.name); + try std.testing.expectEqualStrings("arg1", delta.tool_call.arguments[0].name); } test toArguments { @@ -591,14 +595,6 @@ test toModels { try std.testing.expectEqualStrings("Gemini Model", generic_models[0].display_name); } -test "toDelta with empty arguments_delta" { - const arg_delta = api.InteractionStepDelta{ - .arguments_delta = .{ .arguments = "foo" }, - }; - const delta = toDelta(arg_delta, &.{}); - try std.testing.expect(delta == null); -} - test "toGoogleTool OOM" { const tool = llm.types.Tool{ .name = "get_weather", diff --git a/src/provider/google/provider.zig b/src/provider/google/provider.zig index 42f8fab..ca1beaf 100644 --- a/src/provider/google/provider.zig +++ b/src/provider/google/provider.zig @@ -5,10 +5,13 @@ const makeUri = @import("uri.zig").makeUri; const converter = @import("converter.zig"); const MakeJsonClient = @import("../json_client.zig").MakeJsonClient; const gemini_types = @import("types.zig"); +const accumulator = @import("accumulator.zig"); const Provider = llm.Provider; const ProviderError = llm.Provider.ProviderError; const Allocator = std.mem.Allocator; +const StepAccumulator = accumulator.StepAccumulator; +const TransientDelta = accumulator.TransientDelta; pub const Gemini = MakeProvider(*std.http.Client); @@ -104,10 +107,10 @@ fn MakeProvider(comptime ClientType: type) type { arena_allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, - previous_step: ?llm.types.StepResult, + previous_step: ?llm.types.StepContinuation, stream: bool, ) !struct { uri: std.Uri, payload: api.CreateInteractionRequest } { - const previous_gemini_step: ?*gemini_types.StepResult = if (previous_step) |step| @ptrCast(@alignCast(step.ptr)) else null; + const previous_gemini_step: ?*gemini_types.StepContinuation = if (previous_step) |step| @ptrCast(@alignCast(step.ptr)) else null; var tools: []api.Tool = &.{}; if (session_config.tools.len > 0) { @@ -147,14 +150,22 @@ fn MakeProvider(comptime ClientType: type) type { /// The caller **MUST** call `deinit()` on the returned `StepResult` to free its allocated contents. /// /// Returns the `StepResult` detailing thoughts, outputs, and tool calls, or a `ProviderError` on failure. - pub fn executeStep(ctx: *anyopaque, allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, previous_step: ?llm.types.StepResult) !llm.types.StepResult { + pub fn executeStep( + ctx: *anyopaque, + allocator: Allocator, + session_config: llm.types.SessionConfig, + input: []const llm.types.Step, + previous_step: ?llm.types.StepContinuation, + ) !llm.types.StepOutcome { const self: *Self = @ptrCast(@alignCast(ctx)); var arena = std.heap.ArenaAllocator.init(allocator); defer arena.deinit(); const req_data = try self.buildCreateInteractionRequest(arena.allocator(), session_config, input, previous_step, false); const response = self.rpc_client.postRequest(allocator, api.CreateInteractionRequest, api.Interaction, req_data.uri, req_data.payload) catch return ProviderError.HttpRequestFailed; - return try gemini_types.StepResult.init(allocator, response); + const step_result = try gemini_types.StepResult.init(allocator, response); + const step_continuation = try gemini_types.StepContinuation.init(allocator, response.value.id); + return .{ .result = step_result, .continuation = step_continuation }; } /// Executes a single interaction step using the Gemini API, streaming chunks of the response back. @@ -178,10 +189,10 @@ fn MakeProvider(comptime ClientType: type) type { allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, - previous_step: ?llm.types.StepResult, + previous_step: ?llm.types.StepContinuation, callback: llm.types.StreamingCallback, callback_context: ?*anyopaque, - ) !llm.types.StepResult { + ) !llm.types.StepOutcome { const self: *Self = @ptrCast(@alignCast(ctx)); var arena = std.heap.ArenaAllocator.init(allocator); defer arena.deinit(); @@ -271,31 +282,13 @@ fn MakeProvider(comptime ClientType: type) type { const acc: *StepAccumulator = &step_accumulators.items[e.index]; acc.appendStep(e.delta) catch |err| return Helper.mapError(err); - var arguments: []const llm.types.Argument = &.{}; - defer { - for (arguments) |arg| { - allocator.free(arg.name); - switch (arg.value) { - .string => |s| allocator.free(s), - else => {}, - } - } - allocator.free(arguments); - } - - if (e.delta == .arguments_delta) { - if (try acc.handleToolCallDelta(allocator)) |args| { - arguments = args; - } - } - const delta_payload = converter.toDelta(e.delta, arguments); - - if (delta_payload) |payload| { + if (try acc.buildDelta(allocator, e.delta)) |t_delta| { + defer t_delta.deinit(); callback(callback_context, .{ .event = .{ .step_event = .{ .index = e.index, - .event = .{ .delta = payload }, + .event = .{ .delta = t_delta.delta }, }, }, }); @@ -335,9 +328,18 @@ fn MakeProvider(comptime ClientType: type) type { }); }, .tool_call => |*tc| { - const json_parsed_args = std.json.parseFromSlice(std.json.Value, result_arena.allocator(), tc.arguments_json.written(), .{}) catch return ProviderError.BadResponse; - const function_arguments = api.FunctionArgument.parseFromJsonObject(result_arena.allocator(), json_parsed_args.value) catch return ProviderError.BadResponse; - const parsed_args = converter.toArguments(result_arena.allocator(), function_arguments) catch return ProviderError.BadResponse; + const json_parsed_args = std.json.parseFromSlice(std.json.Value, result_arena.allocator(), tc.arguments_json.written(), .{}) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return ProviderError.BadResponse; + }; + const function_arguments = api.FunctionArgument.parseFromJsonObject(result_arena.allocator(), json_parsed_args.value) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return ProviderError.BadResponse; + }; + const parsed_args = converter.toArguments(result_arena.allocator(), function_arguments) catch |err| { + if (err == error.OutOfMemory) return error.OutOfMemory; + return ProviderError.BadResponse; + }; try final_tool_calls.append(result_arena.allocator(), .{ .id = tc.id, .name = tc.name, @@ -350,158 +352,19 @@ fn MakeProvider(comptime ClientType: type) type { const final_id = interaction_id orelse try result_arena.allocator().dupe(u8, "unknown"); - return try gemini_types.StreamingStepResult.init( + const step_result = try gemini_types.StreamingStepResult.init( allocator, result_arena, - final_id, try final_model_outputs.toOwnedSlice(result_arena.allocator()), try final_thoughts.toOwnedSlice(result_arena.allocator()), try final_tool_calls.toOwnedSlice(result_arena.allocator()), ); + const step_continuation = try gemini_types.StepContinuation.init(allocator, final_id); + return .{ .result = step_result, .continuation = step_continuation }; } }; } -/// Accumulator union used to buffer and aggregate streaming responses for each step. -/// Handles text/thought aggregation and tool call parsing during streaming. -const StepAccumulator = union(enum) { - /// Holds the aggregated thinking/thought process content. - thought: std.Io.Writer.Allocating, - /// Holds the aggregated model response output text. - model_output: std.Io.Writer.Allocating, - /// Holds tool call details and buffers the arguments JSON stream. - tool_call: struct { - /// The unique ID of the tool call. - id: []const u8, - /// The name of the function to invoke. - name: []const u8, - /// Dynamic writer buffer accumulating arguments JSON text. - arguments_json: std.Io.Writer.Allocating, - /// Tracks the count of arguments that have already been streamed out to the callback. - processed_argument_count: usize, - }, - /// Represents an uninitialized or empty accumulator slot. - empty, - - /// Initializes the `StepAccumulator` with an initial step context. - /// - /// `result_arena_allocator` is used to allocate the internal writers and copy initial values. - /// Memory allocated in `result_arena_allocator` is owned by the caller's session arena. - /// `initial_step` provides the starting type and content (e.g. initial thought, model output, or function call). - /// - /// Returns the initialized `StepAccumulator`, or an allocation error. - pub fn init(result_arena_allocator: Allocator, initial_step: api.Step) !@This() { - switch (initial_step) { - .thought => |contents| { - var list = std.Io.Writer.Allocating.init(result_arena_allocator); - for (contents) |c| { - if (c.text) |t| { - try list.writer.writeAll(t); - } - } - return .{ .thought = list }; - }, - .model_output => |contents| { - var list = std.Io.Writer.Allocating.init(result_arena_allocator); - for (contents) |c| { - if (c.text) |t| { - try list.writer.writeAll(t); - } - } - return .{ .model_output = list }; - }, - .function_call => |call| return .{ - .tool_call = .{ - .id = try result_arena_allocator.dupe(u8, call.id), - .name = try result_arena_allocator.dupe(u8, call.name), - .arguments_json = std.Io.Writer.Allocating.init(result_arena_allocator), - .processed_argument_count = 0, - }, - }, - } - } - - /// Appends stream delta data to the accumulator. - /// - /// `step_self` points to the accumulator instance. - /// `delta` is the stream delta payload containing new text or argument fragments. - /// - /// Returns a `ProviderError.BadResponse` if the delta type does not match the active accumulator variant. - pub fn appendStep(step_self: *@This(), delta: api.InteractionStepDelta) !void { - switch (step_self.*) { - .thought => |*list| { - if (delta != .thought_summary) return ProviderError.BadResponse; - if (delta.thought_summary.content.text) |t| { - try list.writer.writeAll(t); - } - }, - .model_output => |*list| { - if (delta != .text_delta) return ProviderError.BadResponse; - if (delta.text_delta.text) |t| { - try list.writer.writeAll(t); - } - }, - .tool_call => |*tc| { - if (delta != .arguments_delta) return ProviderError.BadResponse; - try tc.arguments_json.writer.writeAll(delta.arguments_delta.arguments); - }, - .empty => { - return ProviderError.BadResponse; - }, - } - } - - /// Processes newly accumulated tool call arguments JSON and parses new complete arguments. - /// - /// `acc` is the active `StepAccumulator` (must be the `.tool_call` variant). - /// `allocator` is used to allocate the returned array of arguments. - /// - /// **Memory Alert**: - /// The caller takes ownership of the returned slice of `llm.types.Argument` and **MUST** free - /// the slice, the `name` strings, and the `value` strings when done. - /// - /// Returns a slice of newly completed `llm.types.Argument` parameters, or `null` if no new arguments - /// were completed or parsing failed. - pub fn handleToolCallDelta( - acc: *StepAccumulator, - allocator: Allocator, - ) !?[]const llm.types.Argument { - const json_parsed_args = std.json.parseFromSlice( - std.json.Value, - allocator, - acc.tool_call.arguments_json.written(), - .{}, - ) catch return null; - defer json_parsed_args.deinit(); - - const function_arguments = api.FunctionArgument.parseFromJsonObject( - allocator, - json_parsed_args.value, - ) catch return null; - defer allocator.free(function_arguments); - - const new_count = function_arguments.len - acc.tool_call.processed_argument_count; - if (new_count == 0) { - return null; - } - - const arguments = try converter.dupeArguments(allocator, function_arguments[acc.tool_call.processed_argument_count..]); - errdefer { - for (arguments) |arg| { - allocator.free(arg.name); - switch (arg.value) { - .string => |s| allocator.free(s), - else => {}, - } - } - allocator.free(arguments); - } - - acc.tool_call.processed_argument_count = function_arguments.len; - return arguments; - } -}; - const testing = @import("testing"); test "Gemini initialization and deinitialization" { @@ -646,16 +509,19 @@ test "Gemini.executeStep success" { .{ .prompt = "Hello" }, }; - var result = try p.executeStep(allocator, config, input, null); - defer result.deinit(); + var outcome = try p.executeStep(allocator, config, input, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const result = outcome.result; + const continuation = outcome.continuation; try std.testing.expectEqual(1, result.model_output.len); try std.testing.expectEqualStrings("Hello user!", result.model_output[0].text); try std.testing.expectEqual(1, result.thoughts.len); try std.testing.expectEqualStrings("Thinking...", result.thoughts[0].text); - const gemini_result: *gemini_types.StepResult = @ptrCast(@alignCast(result.ptr)); - try std.testing.expectEqualStrings("interaction_123", gemini_result.interaction_id); + const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + try std.testing.expectEqualStrings("interaction_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } @@ -740,8 +606,11 @@ test "Gemini.executeStep with previous step" { const input1 = &[_]llm.types.Step{ .{ .prompt = "Hello" }, }; - var result1 = try p.executeStep(allocator, config, input1, null); - defer result1.deinit(); + var outcome1 = try p.executeStep(allocator, config, input1, null); + defer outcome1.result.deinit(); + defer outcome1.continuation.deinit(); + const result1 = outcome1.result; + const continuation1 = outcome1.continuation; try std.testing.expectEqualStrings("Hello user!", result1.model_output[0].text); @@ -749,8 +618,10 @@ test "Gemini.executeStep with previous step" { const input2 = &[_]llm.types.Step{ .{ .prompt = "Next prompt" }, }; - var result2 = try p.executeStep(allocator, config, input2, result1); - defer result2.deinit(); + var outcome2 = try p.executeStep(allocator, config, input2, continuation1); + defer outcome2.result.deinit(); + defer outcome2.continuation.deinit(); + const result2 = outcome2.result; try std.testing.expectEqualStrings("Next response", result2.model_output[0].text); try std.testing.expectEqual(1, call_counts[0]); @@ -820,8 +691,10 @@ test "Gemini.executeStep with tools" { .{ .prompt = "Hello" }, }; - var result = try p.executeStep(allocator, config, input, null); - defer result.deinit(); + var outcome = try p.executeStep(allocator, config, input, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const result = outcome.result; try std.testing.expectEqualStrings("Weather is nice!", result.model_output[0].text); try std.testing.expectEqual(1, call_counts[0]); @@ -960,10 +833,10 @@ test "Gemini.executeStepStreaming success" { .thought => |t| { self.thought_text.writer.writeAll(t.text) catch {}; }, - .tool_call => |args| { - if (args.len > 0) { - std.testing.expectEqualStrings("location", args[0].name) catch {}; - if (std.mem.eql(u8, args[0].value.string, "San Francisco")) { + .tool_call => |tc| { + if (tc.arguments.len > 0) { + std.testing.expectEqualStrings("location", tc.arguments[0].name) catch {}; + if (std.mem.eql(u8, tc.arguments[0].value.string, "San Francisco")) { self.tool_call_arguments_received = true; } } @@ -984,8 +857,11 @@ test "Gemini.executeStepStreaming success" { defer ctx.model_output_text.deinit(); defer ctx.thought_text.deinit(); - var result = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); - defer result.deinit(); + var outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const result = outcome.result; + const continuation = outcome.continuation; // Verify callback executions try std.testing.expect(ctx.chunks_received > 0); @@ -1007,8 +883,8 @@ test "Gemini.executeStepStreaming success" { try std.testing.expectEqualStrings("location", result.tool_calls[0].arguments[0].name); try std.testing.expectEqualStrings("San Francisco", result.tool_calls[0].arguments[0].value.string); - const gemini_result: *gemini_types.StreamingStepResult = @ptrCast(@alignCast(result.ptr)); - try std.testing.expectEqualStrings("interaction_streaming_123", gemini_result.interaction_id); + const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + try std.testing.expectEqualStrings("interaction_streaming_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } @@ -1120,10 +996,10 @@ test "Gemini.executeStepStreaming with CRLF line endings" { .thought => |t| { self.thought_text.writer.writeAll(t.text) catch {}; }, - .tool_call => |args| { - if (args.len > 0) { - std.testing.expectEqualStrings("location", args[0].name) catch {}; - if (std.mem.eql(u8, args[0].value.string, "San Francisco")) { + .tool_call => |tc| { + if (tc.arguments.len > 0) { + std.testing.expectEqualStrings("location", tc.arguments[0].name) catch {}; + if (std.mem.eql(u8, tc.arguments[0].value.string, "San Francisco")) { self.tool_call_arguments_received = true; } } @@ -1144,8 +1020,11 @@ test "Gemini.executeStepStreaming with CRLF line endings" { defer ctx.model_output_text.deinit(); defer ctx.thought_text.deinit(); - var result = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); - defer result.deinit(); + var outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const result = outcome.result; + const continuation = outcome.continuation; // Verify callback executions try std.testing.expect(ctx.chunks_received > 0); @@ -1166,103 +1045,11 @@ test "Gemini.executeStepStreaming with CRLF line endings" { try std.testing.expectEqualStrings("location", result.tool_calls[0].arguments[0].name); try std.testing.expectEqualStrings("San Francisco", result.tool_calls[0].arguments[0].value.string); - const gemini_result: *gemini_types.StreamingStepResult = @ptrCast(@alignCast(result.ptr)); - try std.testing.expectEqualStrings("interaction_streaming_123", gemini_result.interaction_id); + const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + try std.testing.expectEqualStrings("interaction_streaming_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } -test "StepAccumulator mismatch delta and init thought" { - const allocator = std.testing.allocator; - - var thoughts = [_]api.Content{.{ .type = .text, .text = "thought_1" }}; - const init_step = api.Step{ .thought = &thoughts }; - var acc = try StepAccumulator.init(allocator, init_step); - defer acc.thought.deinit(); - try std.testing.expect(acc == .thought); - try std.testing.expectEqualStrings("thought_1", acc.thought.written()); - - // Append mismatched delta (should fail) - const mismatch_delta = api.InteractionStepDelta{ - .text_delta = .{ .type = .text, .text = "text" }, - }; - try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); -} - -test "StepAccumulator mismatch delta and init model_output" { - const allocator = std.testing.allocator; - - var outputs = [_]api.Content{.{ .type = .text, .text = "output_1" }}; - const init_step = api.Step{ .model_output = &outputs }; - var acc = try StepAccumulator.init(allocator, init_step); - defer acc.model_output.deinit(); - try std.testing.expect(acc == .model_output); - try std.testing.expectEqualStrings("output_1", acc.model_output.written()); - - // Append mismatched delta (should fail) - const mismatch_delta = api.InteractionStepDelta{ - .arguments_delta = .{ .arguments = "args" }, - }; - try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); -} - -test "StepAccumulator mismatch delta and init tool_call" { - const allocator = std.testing.allocator; - - const init_step = api.Step{ - .function_call = .{ - .id = "fc_id", - .name = "fc_name", - .arguments = &.{}, - }, - }; - var acc = try StepAccumulator.init(allocator, init_step); - defer { - allocator.free(acc.tool_call.id); - allocator.free(acc.tool_call.name); - acc.tool_call.arguments_json.deinit(); - } - try std.testing.expect(acc == .tool_call); - try std.testing.expectEqualStrings("fc_id", acc.tool_call.id); - - // Append mismatched delta (should fail) - const mismatch_delta = api.InteractionStepDelta{ - .thought_summary = .{ .content = .{ .type = .text, .text = "think" } }, - }; - try std.testing.expectError(error.BadResponse, acc.appendStep(mismatch_delta)); - - // Test handleToolCallDelta when json is not valid (returns null) - try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); - - // Test append valid delta and handleToolCallDelta returning null when new_count == 0 - const args_delta = api.InteractionStepDelta{ - .arguments_delta = .{ .arguments = "{\"arg1\": \"val1\"}" }, - }; - try acc.appendStep(args_delta); - const args = try acc.handleToolCallDelta(allocator); - try std.testing.expect(args != null); - defer { - for (args.?) |arg| { - allocator.free(arg.name); - switch (arg.value) { - .string => |s| allocator.free(s), - else => {}, - } - } - allocator.free(args.?); - } - - // Calling handleToolCallDelta again without new changes should return null - try std.testing.expect(try acc.handleToolCallDelta(allocator) == null); -} - -test "StepAccumulator mismatch delta and init empty" { - var acc: StepAccumulator = .empty; - const delta = api.InteractionStepDelta{ - .text_delta = .{ .type = .text, .text = "text" }, - }; - try std.testing.expectError(error.BadResponse, acc.appendStep(delta)); -} - test "Gemini.executeStep returns function call" { const allocator = std.testing.allocator; @@ -1307,10 +1094,13 @@ test "Gemini.executeStep returns function call" { .tools = &.{}, }; - var result = try p.executeStep(allocator, config, &.{}, null); - defer result.deinit(); + var outcome = try p.executeStep(allocator, config, &.{}, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const result = outcome.result; + const continuation = outcome.continuation; - try std.testing.expectEqualStrings("interaction_fc", (@as(*gemini_types.StepResult, @ptrCast(@alignCast(result.ptr)))).interaction_id); + try std.testing.expectEqualStrings("interaction_fc", (@as(*gemini_types.StepContinuation, @ptrCast(@alignCast(continuation.ptr)))).interaction_id); try std.testing.expectEqual(1, result.tool_calls.len); try std.testing.expectEqualStrings("call_abc", result.tool_calls[0].id); try std.testing.expectEqualStrings("get_weather", result.tool_calls[0].name); @@ -1514,9 +1304,11 @@ test "Gemini.executeStepStreaming fallback interaction ID to unknown" { .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, .tools = &.{}, }; - var result = try p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null); - defer result.deinit(); - try std.testing.expectEqualStrings("unknown", (@as(*gemini_types.StreamingStepResult, @ptrCast(@alignCast(result.ptr)))).interaction_id); + var outcome = try p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + const continuation = outcome.continuation; + try std.testing.expectEqualStrings("unknown", (@as(*gemini_types.StepContinuation, @ptrCast(@alignCast(continuation.ptr)))).interaction_id); } test "ListModelsResult.init OOM" { @@ -1538,5 +1330,60 @@ test "StepResult.init OOM" { test "StreamingStepResult.init OOM" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); - try std.testing.expectError(error.OutOfMemory, gemini_types.StreamingStepResult.init(std.testing.failing_allocator, arena, "id", &.{}, &.{}, &.{})); + try std.testing.expectError(error.OutOfMemory, gemini_types.StreamingStepResult.init(std.testing.failing_allocator, arena, &.{}, &.{}, &.{})); +} + +test "StepContinuation.init OOM" { + try std.testing.expectError(error.OutOfMemory, gemini_types.StepContinuation.init(std.testing.failing_allocator, "id")); +} + +test "Gemini.executeStepStreaming malformed stream payload" { + const allocator = std.testing.allocator; + + const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; + const response_body = + \\event: step.delta + \\data: {invalid_json + \\ + ; + + var call_counts = [_]usize{0}; + const mock_client = testing.MockHttpClient{ + .allocator = allocator, + .expectations = &.{ + .{ + .expected_scheme = "https", + .expected_host = "generativelanguage.googleapis.com", + .expected_path = "/v1beta/interactions", + .expected_query = "key=TEST_API_KEY", + .expected_method = .POST, + .expected_payload = expected_payload, + .response_status = .ok, + .response_body = response_body, + }, + }, + .sequential = false, + .call_counts = &call_counts, + }; + + var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); + var p = prov.provider(); + defer p.deinit(); + + const config = llm.types.SessionConfig{ + .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, + .tools = &.{}, + }; + const input = &[_]llm.types.Step{ + .{ .prompt = "Hello" }, + }; + + const CallbackState = struct { + fn callback(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) void { + _ = ctx; + _ = chunk; + } + }; + + try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, input, null, CallbackState.callback, null)); } diff --git a/src/provider/google/types.zig b/src/provider/google/types.zig index 3593a44..2f18bc6 100644 --- a/src/provider/google/types.zig +++ b/src/provider/google/types.zig @@ -47,8 +47,6 @@ pub const ListModelsResult = struct { /// A Gemini-specific implementation of the `llm.types.StepResult` interface. /// Represents the result of an interaction step with the Gemini API. pub const StepResult = struct { - interaction_id: []const u8, - model_output: std.ArrayList(llm.types.ModelOutput), thoughts: std.ArrayList(llm.types.Thought), tool_calls: std.ArrayList(llm.types.ToolCall), @@ -106,7 +104,6 @@ pub const StepResult = struct { const self = try allocator.create(StepResult); errdefer allocator.destroy(self); self.* = .{ - .interaction_id = parsed_response.value.id, .model_output = model_output_list, .thoughts = thoughts_list, .tool_calls = tool_calls_list, @@ -141,8 +138,6 @@ pub const StepResult = struct { /// A Gemini-specific implementation of the `llm.types.StepResult` interface for streaming requests. pub const StreamingStepResult = struct { - interaction_id: []const u8, - model_output: []const llm.types.ModelOutput, thoughts: []const llm.types.Thought, tool_calls: []const llm.types.ToolCall, @@ -154,7 +149,6 @@ pub const StreamingStepResult = struct { pub fn init( allocator: std.mem.Allocator, arena: std.heap.ArenaAllocator, - interaction_id: []const u8, model_output: []const llm.types.ModelOutput, thoughts: []const llm.types.Thought, tool_calls: []const llm.types.ToolCall, @@ -162,7 +156,6 @@ pub const StreamingStepResult = struct { const self = try allocator.create(StreamingStepResult); errdefer allocator.destroy(self); self.* = .{ - .interaction_id = interaction_id, .model_output = model_output, .thoughts = thoughts, .tool_calls = tool_calls, @@ -179,11 +172,47 @@ pub const StreamingStepResult = struct { }; } - /// Frees the resources associated with the result. + /// Frees the resources associated with the StreamingStepResult. pub fn deinit(ctx: *anyopaque) void { const self: *StreamingStepResult = @ptrCast(@alignCast(ctx)); const allocator = self.allocator; self.arena.deinit(); + self.* = undefined; + allocator.destroy(self); + } +}; + +pub const StepContinuation = struct { + interaction_id: []const u8, + + allocator: std.mem.Allocator, + + /// Initializes a new StepContinuation and returns a reference to it via `llm.types.StepContinuation`. + /// + /// This function creates a copy of the interaction_id and stores it in the new StepContinuation. + /// The StepContinuation will deallocate the interaction_id when it is deinitialized. + pub fn init(allocator: std.mem.Allocator, interaction_id: []const u8) !llm.types.StepContinuation { + const self = try allocator.create(StepContinuation); + errdefer allocator.destroy(self); + const dupe_interaction_id = try allocator.dupe(u8, interaction_id); + errdefer allocator.free(dupe_interaction_id); + self.* = .{ + .interaction_id = dupe_interaction_id, + .allocator = allocator, + }; + + return .{ + .ptr = self, + .vtable = &.{ .deinit = deinit }, + }; + } + + /// Frees the resources associated with the StepContinuation. + pub fn deinit(ctx: *anyopaque) void { + const self: *StepContinuation = @ptrCast(@alignCast(ctx)); + const allocator = self.allocator; + allocator.free(self.interaction_id); + self.* = undefined; allocator.destroy(self); } }; diff --git a/src/provider/root.zig b/src/provider/root.zig index 5efcaa3..1e7bd73 100644 --- a/src/provider/root.zig +++ b/src/provider/root.zig @@ -1,4 +1,5 @@ const std = @import("std"); + pub const Gemini = @import("google/provider.zig").Gemini; test { diff --git a/src/root.zig b/src/root.zig index 472c78e..c12b911 100644 --- a/src/root.zig +++ b/src/root.zig @@ -1,13 +1,5 @@ const std = @import("std"); -const llm = @import("llm"); -const provider = @import("provider"); -test "compile llm module" { - _ = llm.Provider; - _ = llm.types.SessionConfig; -} - -test "compile providers module" { - _ = provider.Gemini; - std.testing.refAllDecls(provider); +test { + std.testing.refAllDecls(@This()); } diff --git a/src/testing/MockHttpClient.zig b/src/testing/MockHttpClient.zig index 013a9c3..0f2e86b 100644 --- a/src/testing/MockHttpClient.zig +++ b/src/testing/MockHttpClient.zig @@ -66,7 +66,9 @@ fn verifyExpectation(options: anytype, expectation: RequestExpectation) !void { if (@hasField(@TypeOf(options), "payload")) { try std.testing.expectEqualStrings(expected, options.payload); } else { - return error.ExpectedPayloadButGotNone; + if (@hasField(@TypeOf(options), "response_writer")) { + return error.ExpectedPayloadButGotNone; + } } } else { if (@hasField(@TypeOf(options), "payload")) { @@ -451,3 +453,67 @@ test "MockHttpClient - verification errors unexpected payload" { .payload = "unexpected_payload_data", })); } + +test "MockHttpClient - sequential request flow" { + const allocator = std.testing.allocator; + const uri_post = try std.Uri.parse("https://example.com/post"); + const uri_get = try std.Uri.parse("https://example.com/get"); + + var call_counts = [_]usize{ 0, 0 }; + var client = MockHttpClient{ + .allocator = allocator, + .expectations = &.{ + .{ + .expected_scheme = "https", + .expected_host = "example.com", + .expected_path = "/post", + .expected_query = null, + .expected_method = .POST, + .expected_payload = "payload123", + .response_status = .ok, + .response_body = "response_post", + }, + .{ + .expected_scheme = "https", + .expected_host = "example.com", + .expected_path = "/get", + .expected_query = null, + .expected_method = .GET, + .expected_payload = null, + .response_status = .ok, + .response_body = "response_get", + }, + }, + .sequential = true, + .call_counts = &call_counts, + }; + + // 1. POST request + { + var req = try client.request(.POST, uri_post, .{}); + defer req.deinit(); + try req.sendBodyComplete("payload123"); + var resp = try req.receiveHead(""); + const reader = resp.reader(&[_]u8{}); + var buf: [100]u8 = undefined; + const n = try reader.readSliceShort(&buf); + try std.testing.expectEqualStrings("response_post", buf[0..n]); + try std.testing.expectEqual(@as(usize, 1), call_counts[0]); + } + + // 2. GET request + { + var req = try client.request(.GET, uri_get, .{}); + defer req.deinit(); + try req.sendBodyComplete(""); + var resp = try req.receiveHead(""); + const reader = resp.reader(&[_]u8{}); + var buf: [100]u8 = undefined; + const n = try reader.readSliceShort(&buf); + try std.testing.expectEqualStrings("response_get", buf[0..n]); + try std.testing.expectEqual(@as(usize, 1), call_counts[1]); + } + + // 3. Too many calls error + try std.testing.expectError(error.TooManyCalls, client.request(.GET, uri_get, .{})); +} diff --git a/src/testing/MockProvider.zig b/src/testing/MockProvider.zig index 948bcab..0f52226 100644 --- a/src/testing/MockProvider.zig +++ b/src/testing/MockProvider.zig @@ -3,10 +3,13 @@ const Allocator = std.mem.Allocator; const llm = @import("llm"); const Provider = llm.Provider; const types = llm.types; + const ListModelsResult = types.ListModelsResult; const StepResult = types.StepResult; +const StepContinuation = types.StepContinuation; const SessionConfig = types.SessionConfig; const Step = types.Step; +const StepOutcome = types.StepOutcome; /// A mock implementation of the `llm.Provider` interface for testing. const MockProvider = @This(); @@ -27,12 +30,15 @@ last_session_config: ?SessionConfig = null, /// Stores the input steps from the last `executeStep` call. last_input: ?[]const Step = null, /// Stores the previous step from the last `executeStep` call. -last_previous_step: ?StepResult = null, +last_previous_step: ?StepContinuation = null, /// Optional fixed result to return from `listModels`. list_models_result: ?(Provider.ProviderError!ListModelsResult) = null, /// Optional fixed result to return from `executeStep`. execute_step_result: ?(Provider.ProviderError!StepResult) = null, +/// Optional fixed continuation to return from `executeStep`. +/// Must be set if `execute_step_result` is set. +execute_step_continuation: ?StepContinuation = null, const vtable = Provider.VTable{ .list_models = MockProvider.list_models, @@ -59,6 +65,11 @@ pub const mock_step_vtable = StepResult.VTable{ .deinit = dummyDeinit, }; +/// A dummy virtual table for `StepContinuation` used by the mock. +pub const mock_continuation_vtable = StepContinuation.VTable{ + .deinit = dummyDeinit, +}; + /// A no-op implementation of `deinit`. fn dummyDeinit(ptr: *anyopaque) void { _ = ptr; @@ -85,8 +96,8 @@ fn execute_step( allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, -) Provider.ProviderError!StepResult { + previous_step: ?StepContinuation, +) Provider.ProviderError!StepOutcome { const self: *MockProvider = @ptrCast(@alignCast(ptr)); self.execute_step_calls += 1; self.last_allocator = allocator; @@ -94,14 +105,18 @@ fn execute_step( self.last_input = input; self.last_previous_step = previous_step; if (self.execute_step_result) |res| { - return res; + const ok = try res; + return StepOutcome{ .result = ok, .continuation = self.execute_step_continuation.? }; } - return StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = ptr, - .vtable = &mock_step_vtable, + return StepOutcome{ + .result = StepResult{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = ptr, + .vtable = &mock_step_vtable, + }, + .continuation = StepContinuation{ .ptr = ptr, .vtable = &mock_continuation_vtable }, }; } @@ -111,10 +126,10 @@ fn execute_step_streaming( allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepResult, + previous_step: ?StepContinuation, callback: types.StreamingCallback, callback_context: ?*anyopaque, -) Provider.ProviderError!StepResult { +) Provider.ProviderError!StepOutcome { const self: *MockProvider = @ptrCast(@alignCast(ptr)); _ = callback; _ = callback_context; @@ -124,14 +139,18 @@ fn execute_step_streaming( self.last_input = input; self.last_previous_step = previous_step; if (self.execute_step_result) |res| { - return res; + const ok = try res; + return StepOutcome{ .result = ok, .continuation = self.execute_step_continuation.? }; } - return StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = ptr, - .vtable = &mock_step_vtable, + return StepOutcome{ + .result = StepResult{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &.{}, + .ptr = ptr, + .vtable = &mock_step_vtable, + }, + .continuation = StepContinuation{ .ptr = ptr, .vtable = &mock_continuation_vtable }, }; } @@ -140,3 +159,14 @@ fn deinit(ptr: *anyopaque) void { const self: *MockProvider = @ptrCast(@alignCast(ptr)); self.deinit_calls += 1; } + +pub const MockStepContinuation = struct { + pub fn stepContinuation(self: *MockStepContinuation) StepContinuation { + return .{ + .ptr = self, + .vtable = &.{ + .deinit = dummyDeinit, + }, + }; + } +};