556 lines
19 KiB
C++
556 lines
19 KiB
C++
// Copyright (c) 2025 YAMS Contributors
|
|
// SPDX-License-Identifier: GPL-3.0-or-later
|
|
|
|
// Unit tests for ABI Model Provider dimension handling and error messages
|
|
#include <catch2/catch_test_macros.hpp>
|
|
#include <catch2/matchers/catch_matchers_string.hpp>
|
|
|
|
#include <yams/core/types.h>
|
|
#include <yams/daemon/components/ConfigResolver.h>
|
|
#include <yams/daemon/daemon.h>
|
|
#include <yams/daemon/resource/abi_model_provider_adapter.h>
|
|
#include <yams/plugins/model_provider_v1.h>
|
|
|
|
#include <cstdlib>
|
|
#include <fstream>
|
|
#include <memory>
|
|
#include <random>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
using namespace yams::daemon;
|
|
using namespace yams;
|
|
using Catch::Matchers::ContainsSubstring;
|
|
|
|
namespace {
|
|
|
|
// Mock ONNX plugin provider for testing
|
|
struct MockProviderContext {
|
|
yams_status_t nextStatus = YAMS_OK;
|
|
std::string nextError;
|
|
size_t outputDim = 384;
|
|
std::vector<float> nextEmbedding;
|
|
bool simulateDimensionMismatch = false;
|
|
size_t maxOutputBatch = 0; // 0 = no cap
|
|
bool unknownDimOnFirstBatch = false;
|
|
size_t batchCalls = 0;
|
|
};
|
|
|
|
// Mock implementation of plugin functions
|
|
yams_status_t mock_generate_embedding(void* self, const char* model_id, const uint8_t* text,
|
|
size_t text_len, float** out_vec, size_t* out_dim) {
|
|
(void)model_id;
|
|
(void)text;
|
|
(void)text_len;
|
|
auto* ctx = static_cast<MockProviderContext*>(self);
|
|
if (!ctx || !out_vec || !out_dim) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
if (ctx->simulateDimensionMismatch) {
|
|
// Return wrong dimension
|
|
*out_dim = 768; // Wrong! Config expects 384
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
if (ctx->nextStatus != YAMS_OK) {
|
|
return ctx->nextStatus;
|
|
}
|
|
// Allocate embedding with configured dimension
|
|
*out_dim = ctx->outputDim;
|
|
*out_vec = static_cast<float*>(malloc(*out_dim * sizeof(float)));
|
|
for (size_t i = 0; i < *out_dim; ++i) {
|
|
(*out_vec)[i] = static_cast<float>(i) * 0.01f;
|
|
}
|
|
return YAMS_OK;
|
|
}
|
|
|
|
yams_status_t mock_generate_embedding_batch(void* self, const char* model_id,
|
|
const uint8_t* const* inputs, const size_t* input_lens,
|
|
size_t batch_size, float** out_vecs, size_t* out_batch,
|
|
size_t* out_dim) {
|
|
(void)model_id;
|
|
auto* ctx = static_cast<MockProviderContext*>(self);
|
|
if (!ctx || !out_vecs || !out_batch || !out_dim) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
ctx->batchCalls++;
|
|
if (ctx->simulateDimensionMismatch) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
if (ctx->nextStatus != YAMS_OK) {
|
|
return ctx->nextStatus;
|
|
}
|
|
if (!inputs || !input_lens || batch_size == 0) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
// Allocate batch embeddings
|
|
if (ctx->maxOutputBatch > 0 && batch_size > ctx->maxOutputBatch) {
|
|
*out_batch = ctx->maxOutputBatch;
|
|
} else {
|
|
*out_batch = batch_size;
|
|
}
|
|
if (ctx->unknownDimOnFirstBatch && ctx->batchCalls == 1) {
|
|
*out_dim = 0;
|
|
*out_vecs = nullptr;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
*out_dim = ctx->outputDim;
|
|
*out_vecs = static_cast<float*>(malloc((*out_batch) * ctx->outputDim * sizeof(float)));
|
|
for (size_t i = 0; i < (*out_batch) * ctx->outputDim; ++i) {
|
|
(*out_vecs)[i] = static_cast<float>(i % ctx->outputDim) * 0.01f;
|
|
}
|
|
return YAMS_OK;
|
|
}
|
|
|
|
void mock_free_embedding(void* self, float* vec, size_t dim) {
|
|
(void)self;
|
|
(void)dim;
|
|
free(vec);
|
|
}
|
|
|
|
void mock_free_embedding_batch(void* self, float* vecs, size_t batch, size_t dim) {
|
|
(void)self;
|
|
(void)batch;
|
|
(void)dim;
|
|
free(vecs);
|
|
}
|
|
|
|
yams_status_t mock_load_model(void* self, const char* model_id, const char* model_path,
|
|
const char* options) {
|
|
(void)model_id;
|
|
(void)model_path;
|
|
(void)options;
|
|
auto* ctx = static_cast<MockProviderContext*>(self);
|
|
return ctx->nextStatus;
|
|
}
|
|
|
|
yams_status_t mock_is_model_loaded(void* self, const char* model_id, bool* loaded) {
|
|
(void)self;
|
|
(void)model_id;
|
|
if (!loaded) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
*loaded = true;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
yams_status_t mock_get_embedding_dim(void* self, const char* model_id, size_t* dim) {
|
|
(void)model_id;
|
|
auto* ctx = static_cast<MockProviderContext*>(self);
|
|
if (!ctx || !dim) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
*dim = ctx->outputDim;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
yams_status_t mock_get_loaded_models(void* self, const char*** model_ids, size_t* count) {
|
|
(void)self;
|
|
if (!model_ids || !count) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
static const char* models[] = {"test-model"};
|
|
*model_ids = models;
|
|
*count = 1;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
void mock_free_model_list(void* self, const char** model_ids, size_t count) {
|
|
(void)self;
|
|
(void)model_ids;
|
|
(void)count;
|
|
// Static array, no need to free
|
|
}
|
|
|
|
} // anonymous namespace
|
|
|
|
struct AbiModelProviderDimensionFixture {
|
|
AbiModelProviderDimensionFixture() {
|
|
// Initialize mock context
|
|
mockContext_ = std::make_unique<MockProviderContext>();
|
|
|
|
// Setup provider table
|
|
table_.self = mockContext_.get();
|
|
table_.generate_embedding = mock_generate_embedding;
|
|
table_.generate_embedding_batch = mock_generate_embedding_batch;
|
|
table_.free_embedding = mock_free_embedding;
|
|
table_.free_embedding_batch = mock_free_embedding_batch;
|
|
table_.load_model = mock_load_model;
|
|
table_.is_model_loaded = mock_is_model_loaded;
|
|
table_.get_embedding_dim = mock_get_embedding_dim;
|
|
table_.get_loaded_models = mock_get_loaded_models;
|
|
table_.free_model_list = mock_free_model_list;
|
|
table_.unload_model = nullptr;
|
|
table_.set_progress_callback = nullptr;
|
|
table_.get_runtime_info_json = nullptr;
|
|
table_.free_string = nullptr;
|
|
|
|
adapter_ = std::make_unique<AbiModelProviderAdapter>(&table_);
|
|
}
|
|
|
|
~AbiModelProviderDimensionFixture() {
|
|
adapter_.reset();
|
|
mockContext_.reset();
|
|
}
|
|
|
|
std::unique_ptr<MockProviderContext> mockContext_;
|
|
yams_model_provider_v1 table_{};
|
|
std::unique_ptr<AbiModelProviderAdapter> adapter_;
|
|
};
|
|
|
|
// ============================================================================
|
|
// Test 1: Error Message Preservation
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: error code mapping",
|
|
"[daemon]") {
|
|
SECTION("Success maps correctly") {
|
|
mockContext_->nextStatus = YAMS_OK;
|
|
auto result = adapter_->loadModel("test-model");
|
|
CHECK(result);
|
|
}
|
|
|
|
SECTION("Invalid argument maps correctly") {
|
|
mockContext_->nextStatus = YAMS_ERR_INVALID_ARG;
|
|
auto result = adapter_->loadModel("test-model");
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::InvalidArgument);
|
|
}
|
|
|
|
SECTION("Not found maps correctly") {
|
|
mockContext_->nextStatus = YAMS_ERR_NOT_FOUND;
|
|
auto result = adapter_->loadModel("test-model");
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::NotFound);
|
|
}
|
|
|
|
SECTION("IO error maps correctly") {
|
|
mockContext_->nextStatus = YAMS_ERR_IO;
|
|
auto result = adapter_->loadModel("test-model");
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::IOError);
|
|
}
|
|
|
|
SECTION("Internal error maps correctly") {
|
|
mockContext_->nextStatus = YAMS_ERR_INTERNAL;
|
|
auto result = adapter_->loadModel("test-model");
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::InternalError);
|
|
}
|
|
|
|
SECTION("Unsupported maps correctly") {
|
|
mockContext_->nextStatus = YAMS_ERR_UNSUPPORTED;
|
|
auto result = adapter_->loadModel("test-model");
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::NotImplemented);
|
|
}
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: error context includes operation", "[daemon]") {
|
|
mockContext_->nextStatus = YAMS_ERR_INVALID_ARG;
|
|
|
|
auto result = adapter_->generateEmbedding("test text");
|
|
|
|
REQUIRE_FALSE(result);
|
|
CHECK_THAT(result.error().message, ContainsSubstring("generate_embedding"));
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: batch error context includes operation", "[daemon]") {
|
|
mockContext_->nextStatus = YAMS_ERR_INVALID_ARG;
|
|
|
|
std::vector<std::string> texts = {"text1", "text2"};
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE_FALSE(result);
|
|
CHECK_THAT(result.error().message, ContainsSubstring("generate_embedding_batch"));
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 2: Dimension Detection
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: get embedding dim 384 (MiniLM-like)", "[daemon]") {
|
|
mockContext_->outputDim = 384;
|
|
|
|
size_t dim = adapter_->getEmbeddingDim("test-model");
|
|
|
|
CHECK(dim == 384);
|
|
}
|
|
|
|
TEST_CASE("ConfigResolver prefers embeddinggemma from preload list", "[daemon][config]") {
|
|
namespace fs = std::filesystem;
|
|
|
|
const fs::path configPath =
|
|
fs::temp_directory_path() / "yams_config_resolver_embeddinggemma_test.toml";
|
|
{
|
|
std::ofstream out(configPath, std::ios::trunc);
|
|
REQUIRE(out.is_open());
|
|
out << "[daemon.models]\n";
|
|
out << "preload_models = [\"embeddinggemma-300m\"]\n";
|
|
}
|
|
|
|
yams::daemon::DaemonConfig cfg;
|
|
cfg.configFilePath = configPath;
|
|
|
|
const auto preferred = yams::daemon::ConfigResolver::resolvePreferredModel(cfg, {});
|
|
CHECK(preferred == "embeddinggemma-300m");
|
|
|
|
std::error_code ec;
|
|
fs::remove(configPath, ec);
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: get embedding dim 768 (MPNet-like)", "[daemon]") {
|
|
mockContext_->outputDim = 768;
|
|
|
|
size_t dim = adapter_->getEmbeddingDim("test-model");
|
|
|
|
CHECK(dim == 768);
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: get embedding dim 1024 (E5-large-like)", "[daemon]") {
|
|
mockContext_->outputDim = 1024;
|
|
|
|
size_t dim = adapter_->getEmbeddingDim("test-model");
|
|
|
|
CHECK(dim == 1024);
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: generated embedding matches declared dim", "[daemon]") {
|
|
mockContext_->outputDim = 384;
|
|
|
|
auto result = adapter_->generateEmbedding("test text");
|
|
|
|
REQUIRE(result);
|
|
CHECK(result.value().size() == 384);
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: batch embeddings match declared dim", "[daemon]") {
|
|
mockContext_->outputDim = 768;
|
|
|
|
std::vector<std::string> texts = {"text1", "text2", "text3"};
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE(result);
|
|
REQUIRE(result.value().size() == 3);
|
|
for (size_t i = 0; i < result.value().size(); ++i) {
|
|
INFO("Checking embedding " << i);
|
|
CHECK(result.value()[i].size() == 768);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 3: Model Info Includes Dimension
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: model info includes dimension", "[daemon]") {
|
|
mockContext_->outputDim = 384;
|
|
mockContext_->nextStatus = YAMS_OK;
|
|
|
|
// First load the model
|
|
auto loadResult = adapter_->loadModel("test-model");
|
|
REQUIRE(loadResult);
|
|
|
|
auto infoResult = adapter_->getModelInfo("test-model");
|
|
|
|
REQUIRE(infoResult);
|
|
CHECK(infoResult.value().embeddingDim == 384);
|
|
CHECK(infoResult.value().name == "test-model");
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 4: Dimension Mismatch Error Scenarios
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: dimension mismatch in generate", "[daemon]") {
|
|
mockContext_->simulateDimensionMismatch = true;
|
|
|
|
auto result = adapter_->generateEmbedding("test text");
|
|
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::InvalidArgument);
|
|
// TODO: After enhancement, verify error includes dimension details
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: dimension mismatch in batch",
|
|
"[daemon]") {
|
|
mockContext_->simulateDimensionMismatch = true;
|
|
|
|
std::vector<std::string> texts = {"text1", "text2"};
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::InvalidArgument);
|
|
// TODO: After enhancement, verify error includes dimension details
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 5: Empty/Invalid Input Handling
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: empty text generate",
|
|
"[daemon]") {
|
|
auto result = adapter_->generateEmbedding("");
|
|
|
|
// Should still work - empty text is valid input
|
|
CHECK(result);
|
|
if (result) {
|
|
CHECK(result.value().size() == mockContext_->outputDim);
|
|
}
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: empty batch", "[daemon]") {
|
|
std::vector<std::string> emptyBatch;
|
|
|
|
auto result = adapter_->generateBatchEmbeddings(emptyBatch);
|
|
|
|
// Empty batch should return empty result, not error
|
|
REQUIRE(result);
|
|
CHECK(result.value().empty());
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: batch with empty strings",
|
|
"[daemon]") {
|
|
std::vector<std::string> texts = {"", "", ""};
|
|
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE(result);
|
|
CHECK(result.value().size() == 3);
|
|
for (const auto& emb : result.value()) {
|
|
CHECK(emb.size() == mockContext_->outputDim);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 6: Multiple Models with Different Dimensions
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: switch between dimensions",
|
|
"[daemon]") {
|
|
SECTION("384-dim model") {
|
|
mockContext_->outputDim = 384;
|
|
auto result384 = adapter_->generateEmbedding("text");
|
|
REQUIRE(result384);
|
|
CHECK(result384.value().size() == 384);
|
|
}
|
|
|
|
SECTION("768-dim model") {
|
|
mockContext_->outputDim = 768;
|
|
auto result768 = adapter_->generateEmbedding("text");
|
|
REQUIRE(result768);
|
|
CHECK(result768.value().size() == 768);
|
|
}
|
|
|
|
SECTION("1024-dim model") {
|
|
mockContext_->outputDim = 1024;
|
|
auto result1024 = adapter_->generateEmbedding("text");
|
|
REQUIRE(result1024);
|
|
CHECK(result1024.value().size() == 1024);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Test 7: Large Batch Performance Validation
|
|
// ============================================================================
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture, "AbiModelProvider: large batch dimensions",
|
|
"[daemon]") {
|
|
mockContext_->outputDim = 384;
|
|
|
|
// Create batch of 100 texts
|
|
std::vector<std::string> texts;
|
|
for (int i = 0; i < 100; ++i) {
|
|
texts.push_back("Test text " + std::to_string(i));
|
|
}
|
|
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE(result);
|
|
CHECK(result.value().size() == 100);
|
|
for (const auto& emb : result.value()) {
|
|
CHECK(emb.size() == 384);
|
|
}
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: provider partial batch output handled safely",
|
|
"[daemon][abi][regression]") {
|
|
mockContext_->outputDim = 768;
|
|
mockContext_->maxOutputBatch = 2;
|
|
|
|
std::vector<std::string> texts = {"text1", "text2", "text3", "text4"};
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::ResourceExhausted);
|
|
CHECK_THAT(result.error().message, ContainsSubstring("generate_embedding_batch_partial"));
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: deterministic fuzz-style partial batch invariants",
|
|
"[daemon][abi][regression][fuzz-lite]") {
|
|
std::mt19937 rng(0x5A17BEEF);
|
|
std::uniform_int_distribution<size_t> dimDist(1, 1024);
|
|
std::uniform_int_distribution<size_t> batchDist(1, 64);
|
|
|
|
for (int iter = 0; iter < 100; ++iter) {
|
|
const size_t requestedBatch = batchDist(rng);
|
|
const size_t outputDim = dimDist(rng);
|
|
std::uniform_int_distribution<size_t> capDist(1, requestedBatch);
|
|
const size_t cap = capDist(rng);
|
|
|
|
mockContext_->outputDim = outputDim;
|
|
mockContext_->maxOutputBatch = cap;
|
|
|
|
std::vector<std::string> texts;
|
|
texts.reserve(requestedBatch);
|
|
for (size_t i = 0; i < requestedBatch; ++i) {
|
|
texts.push_back("seeded-text-" + std::to_string(iter) + "-" + std::to_string(i));
|
|
}
|
|
|
|
auto result = adapter_->generateBatchEmbeddings(texts);
|
|
INFO("iter=" << iter << " requestedBatch=" << requestedBatch << " cap=" << cap
|
|
<< " outputDim=" << outputDim);
|
|
if (cap < requestedBatch) {
|
|
REQUIRE_FALSE(result);
|
|
CHECK(result.error().code == yams::ErrorCode::ResourceExhausted);
|
|
CHECK_THAT(result.error().message,
|
|
ContainsSubstring("generate_embedding_batch_partial"));
|
|
} else {
|
|
REQUIRE(result);
|
|
REQUIRE(result.value().size() == requestedBatch);
|
|
for (const auto& emb : result.value()) {
|
|
CHECK(emb.size() == outputDim);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
TEST_CASE_METHOD(AbiModelProviderDimensionFixture,
|
|
"AbiModelProvider: unknown initial dim with truncated batch stays safe",
|
|
"[daemon][abi][regression]") {
|
|
mockContext_->outputDim = 768;
|
|
mockContext_->maxOutputBatch = 2;
|
|
mockContext_->unknownDimOnFirstBatch = true;
|
|
|
|
std::vector<std::string> texts = {"text1", "text2", "text3", "text4"};
|
|
|
|
auto first = adapter_->generateBatchEmbeddings(texts);
|
|
REQUIRE_FALSE(first);
|
|
CHECK(first.error().code == yams::ErrorCode::ResourceExhausted);
|
|
CHECK_THAT(first.error().message, ContainsSubstring("generate_embedding_batch_partial"));
|
|
|
|
auto second = adapter_->generateBatchEmbeddings(texts);
|
|
REQUIRE_FALSE(second);
|
|
CHECK(second.error().code == yams::ErrorCode::ResourceExhausted);
|
|
CHECK_THAT(second.error().message, ContainsSubstring("generate_embedding_batch_partial"));
|
|
}
|