#include <atomic>
#include <chrono>
#include <cstddef>
#include <span>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
#include <agent_smith/llm_client.hpp>
#include <agent_smith/sync_wait.hpp>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mw/http_client_async_mock.hpp>
namespace agent_smith
{
namespace
{
using ::testing::_;
using ::testing::InSequence;
using ::testing::Invoke;
mw::URL testEndpoint()
{
auto endpoint = mw::URL::fromStr("https://llm.example/api");
EXPECT_TRUE(endpoint.has_value());
return std::move(*endpoint);
}
LlmClientConfig testConfig()
{
return LlmClientConfig{
.endpoint = testEndpoint(),
.model = "test-model",
.max_retries = 0,
};
}
bool sendChunk(const mw::ChunkCallback& callback, std::string_view text)
{
return callback(std::as_bytes(std::span(text.data(), text.size())));
}
mw::Task<mw::E<mw::HTTPResponse>>
streamChunks(mw::ChunkCallback callback, Runtime& runtime,
std::vector<std::string> chunks,
std::atomic<bool>& second_chunk_sent)
{
if(chunks.empty() || !sendChunk(callback, chunks.front()))
{
co_return std::unexpected(mw::runtimeError("stream callback aborted"));
}
auto slept = co_await runtime.sleepFor(std::chrono::milliseconds(10));
if(!slept)
{
co_return std::unexpected(slept.error());
}
second_chunk_sent.store(true);
for(std::size_t index = 1; index < chunks.size(); ++index)
{
if(!sendChunk(callback, chunks[index]))
{
co_return std::unexpected(
mw::runtimeError("stream callback aborted"));
}
}
co_return mw::HTTPResponse(200, "");
}
TEST(LlmClientTest, SendsOpenAIRequestAndParsesAssistantResponse)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClientConfig config = testConfig();
config.api_key = "secret";
config.default_sampling.temperature = 0.2;
config.default_sampling.extra = {{"provider_option", true}};
LlmClient client(runtime, http, std::move(config));
mw::HTTPRequest captured;
EXPECT_CALL(http, post(_))
.WillOnce(Invoke(
[&](mw::HTTPRequest request)
{
captured = std::move(request);
return mw::HTTPSessionAsyncMock::complete(mw::HTTPResponse(
200,
R"({"choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}})"));
}));
LlmRequest request{
.messages = {SystemMessage{"be concise"}, UserMessage{"hi"}},
.tools =
{
ToolSchema{
.name = "weather",
.description = "Look up the weather.",
.parameters =
{
{"type", "object"},
{"properties", {{"city", {{"type", "string"}}}}},
},
},
},
.sampling =
SamplingParams{
.top_p = 0.9,
.max_tokens = 32,
},
};
auto result = syncWait(client.complete(std::move(request)));
ASSERT_TRUE(result.has_value());
EXPECT_EQ(result->message.content, "hello");
ASSERT_TRUE(result->message.usage.has_value());
EXPECT_EQ(result->message.usage->total, 5U);
EXPECT_EQ(captured.url, "https://llm.example/v1/chat/completions");
EXPECT_EQ(captured.header.at("Content-Type"), "application/json");
EXPECT_EQ(captured.header.at("Authorization"), "Bearer secret");
const nlohmann::json body = nlohmann::json::parse(captured.request_data);
EXPECT_EQ(body.at("model"), "test-model");
EXPECT_EQ(body.at("messages"),
nlohmann::json({{{"role", "system"}, {"content", "be concise"}},
{{"role", "user"}, {"content", "hi"}}}));
EXPECT_FALSE(body.at("stream"));
EXPECT_EQ(body.at("temperature"), 0.2);
EXPECT_EQ(body.at("top_p"), 0.9);
EXPECT_EQ(body.at("max_tokens"), 32);
EXPECT_TRUE(body.at("provider_option"));
EXPECT_EQ(body.at("tool_choice"), "auto");
nlohmann::json expected_tools = nlohmann::json::array();
expected_tools.push_back({
{"type", "function"},
{"function",
{
{"name", "weather"},
{"description", "Look up the weather."},
{"parameters",
{
{"type", "object"},
{"properties", {{"city", {{"type", "string"}}}}},
}},
}},
});
EXPECT_EQ(body.at("tools"), expected_tools);
}
TEST(LlmClientTest, RejectsReservedExtraBeforeSendingRequest)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClient client(runtime, http, testConfig());
EXPECT_CALL(http, post(_)).Times(0);
LlmRequest request{
.sampling =
SamplingParams{
.extra = {{"stream", true}},
},
};
auto result = syncWait(client.complete(std::move(request)));
ASSERT_FALSE(result.has_value());
EXPECT_NE(result.error().as<ParseError>(), nullptr);
}
TEST(LlmClientTest, ReturnsTypedErrorForHttpFailure)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClient client(runtime, http, testConfig());
EXPECT_CALL(http, post(_))
.WillOnce(Invoke(
[](mw::HTTPRequest)
{
return mw::HTTPSessionAsyncMock::complete(
mw::HTTPResponse(400, "bad request"));
}));
auto result = syncWait(client.complete({}));
ASSERT_FALSE(result.has_value());
const LlmApiError* error = result.error().as<LlmApiError>();
ASSERT_NE(error, nullptr);
EXPECT_EQ(error->http_status, 400);
EXPECT_EQ(error->body_excerpt, "bad request");
}
TEST(LlmClientTest, ReturnsTypedErrorForMalformedResponse)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClient client(runtime, http, testConfig());
EXPECT_CALL(http, post(_))
.WillOnce(Invoke(
[](mw::HTTPRequest)
{
return mw::HTTPSessionAsyncMock::complete(
mw::HTTPResponse(200, "not json"));
}));
auto result = syncWait(client.complete({}));
ASSERT_FALSE(result.has_value());
EXPECT_NE(result.error().as<ParseError>(), nullptr);
}
TEST(LlmClientTest, RetriesRetryableHttpFailures)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClientConfig config = testConfig();
config.max_retries = 1;
LlmClient client(runtime, http, std::move(config));
{
InSequence sequence;
EXPECT_CALL(http, post(_))
.WillOnce(Invoke(
[](mw::HTTPRequest)
{
return mw::HTTPSessionAsyncMock::complete(
mw::HTTPResponse(503, "temporarily unavailable"));
}));
EXPECT_CALL(http, post(_))
.WillOnce(Invoke(
[](mw::HTTPRequest)
{
return mw::HTTPSessionAsyncMock::complete(mw::HTTPResponse(
200,
R"({"choices":[{"message":{"role":"assistant","content":"recovered"}}]})"));
}));
}
auto result = syncWait(client.complete({}));
ASSERT_TRUE(result.has_value());
EXPECT_EQ(result->message.content, "recovered");
}
TEST(LlmClientTest, StreamsEventsWhileHttpRequestIsActive)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClient client(runtime, http, testConfig());
std::atomic<bool> second_chunk_sent = false;
mw::HTTPRequest captured;
const std::vector<std::string> chunks{
"data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"},"
"\"finish_reason\":null}]}\n\n",
R"(data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"weather","arguments":"{\"city\":"}}]},"finish_reason":null}]}
data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"Portland\"}"}}]},"finish_reason":"tool_calls"}]}
data: {"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}
data: [DONE]
)",
};
EXPECT_CALL(http, postStream(_, _))
.WillOnce(Invoke(
[&](mw::HTTPRequest request, mw::ChunkCallback callback)
{
captured = std::move(request);
return streamChunks(std::move(callback), runtime, chunks,
second_chunk_sent);
}));
std::vector<StreamEvent> events;
bool first_hook_before_second_chunk = false;
auto result = syncWait(client.completeStream(
{},
[&](const StreamEvent& event)
{
if(std::holds_alternative<StreamTokenDelta>(event))
{
first_hook_before_second_chunk = !second_chunk_sent.load();
}
events.push_back(event);
}));
ASSERT_TRUE(result.has_value());
EXPECT_TRUE(first_hook_before_second_chunk);
const nlohmann::json body = nlohmann::json::parse(captured.request_data);
EXPECT_TRUE(body.at("stream"));
EXPECT_FALSE(body.contains("tools"));
ASSERT_EQ(events.size(), 4U);
const auto* token = std::get_if<StreamTokenDelta>(&events[0]);
ASSERT_NE(token, nullptr);
EXPECT_EQ(token->text, "Hello ");
const auto* started = std::get_if<StreamToolCallStarted>(&events[1]);
ASSERT_NE(started, nullptr);
EXPECT_EQ(started->id, "call_1");
EXPECT_EQ(started->name, "weather");
const auto* completed = std::get_if<StreamToolCallCompleted>(&events[2]);
ASSERT_NE(completed, nullptr);
EXPECT_EQ(completed->arguments, nlohmann::json({{"city", "Portland"}}));
const auto* finished = std::get_if<StreamFinished>(&events[3]);
ASSERT_NE(finished, nullptr);
EXPECT_EQ(finished->finish_reason, "tool_calls");
ASSERT_TRUE(finished->usage.has_value());
EXPECT_EQ(finished->usage->total, 5U);
EXPECT_EQ(result->message.content, "Hello ");
ASSERT_EQ(result->message.tool_calls.size(), 1U);
EXPECT_EQ(result->message.tool_calls[0].id, "call_1");
EXPECT_EQ(result->message.tool_calls[0].arguments,
nlohmann::json({{"city", "Portland"}}));
}
TEST(LlmClientTest, RejectsMalformedStreamToolArguments)
{
Runtime runtime;
mw::HTTPSessionAsyncMock http;
LlmClient client(runtime, http, testConfig());
std::atomic<bool> second_chunk_sent = false;
EXPECT_CALL(http, postStream(_, _))
.WillOnce(Invoke(
[&](mw::HTTPRequest, mw::ChunkCallback callback)
{
return streamChunks(
std::move(callback), runtime,
{R"(data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"weather","arguments":"not JSON"}}]},"finish_reason":"tool_calls"}]}
data: [DONE]
)"},
second_chunk_sent);
}));
auto result =
syncWait(client.completeStream({}, [](const StreamEvent&) {}));
ASSERT_FALSE(result.has_value());
EXPECT_NE(result.error().as<ParseError>(), nullptr);
}
} // namespace
} // namespace agent_smith