Set default batch size to 16

src/functions/aggregate/llm_first_or_last/implementation.cpp

@@ -118,7 +118,7 @@ nlohmann::json LlmFirstOrLast::Evaluate(nlohmann::json& tuples) {
             auto result_idx = GetFirstOrLastTupleId(batch_tuples);
 
             batch_tuples.clear();
-            for (auto i = 0; i < static_cast<int>(tuples.size()) - 1; i++) {
+            for (auto i = 0; i < static_cast<int>(tuples.size()); i++) {
                 batch_tuples.push_back(nlohmann::json::object());
                 for (const auto& item: tuples[i].items()) {
                     if (item.key() == "data") {
@@ -139,7 +139,15 @@ nlohmann::json LlmFirstOrLast::Evaluate(nlohmann::json& tuples) {
 
     } while (start_index < static_cast<int>(tuples[0]["data"].size()));
 
-    return batch_tuples;
+    auto result_tuples = nlohmann::json::array();
+    for (const auto& column: batch_tuples) {
+        if (column.contains("name") && column["name"].is_string() &&
+            column["name"].get<std::string>() == "flock_row_id") {
+            continue;
+        }
+        result_tuples.push_back(column);
+    }
+    return result_tuples;
 }
 
 void LlmFirstOrLast::FinalizeResults(duckdb::Vector& states, duckdb::AggregateInputData& aggr_input_data,

src/functions/aggregate/llm_rerank/implementation.cpp

@@ -109,10 +109,7 @@ nlohmann::json LlmRerank::SlidingWindow(nlohmann::json& tuples) {
     auto carry_forward_tuples = nlohmann::json::array();
     int start_index = 0;
 
-    auto batch_size = static_cast<int>(model.GetModelDetails().batch_size);
-    if (batch_size == 2048) {
-        batch_size = std::min<int>(batch_size, num_tuples);
-    }
+    auto batch_size = std::min<int>(model.GetModelDetails().batch_size, num_tuples);
 
     if (batch_size <= 0) {
         throw std::runtime_error("Batch size must be greater than zero");

src/include/flock/model_manager/repository.hpp

@@ -8,6 +8,8 @@
 
 namespace flock {
 
+inline constexpr int DEFAULT_BATCH_SIZE = 16;
+
 struct ModelDetails {
     std::string provider_name;
     std::string model_name;

src/model_manager/model.cpp

@@ -87,7 +87,7 @@ void Model::LoadModelDetails(const nlohmann::json& model_json) {
         } else if (db_model_args.contains("batch_size")) {
             model_details_.batch_size = db_model_args.at("batch_size").get<int>();
         } else {
-            model_details_.batch_size = 2048;
+            model_details_.batch_size = DEFAULT_BATCH_SIZE;
         }
     }
 }

test/unit/functions/aggregate/llm_first.cpp

@@ -82,6 +82,32 @@ TEST_F(LLMFirstTest, MultipleTuplesWithoutGroupBy) {
     ASSERT_EQ(results->GetValue(0, 0).GetValue<std::string>(), GetExpectedResponse());
 }
 
+TEST_F(LLMFirstTest, DefaultBatchSizeSplitsLargeInput) {
+    constexpr size_t input_count = DEFAULT_BATCH_SIZE + 1;
+
+    EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, ::testing::_, ::testing::_, ::testing::_))
+            .Times(2);
+    EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+            .Times(2)
+            .WillRepeatedly(::testing::Return(std::vector<nlohmann::json>{GetExpectedJsonResponse()}));
+
+    auto con = Config::GetConnection();
+
+    const auto results = con.Query(
+            "SELECT llm_first("
+            "{'model_name': 'gpt-4o'}, "
+            "{'prompt': 'What is the most relevant product?', 'context_columns': [{'data': description}]}"
+            ") AS first_product FROM range(" +
+            std::to_string(input_count) + ") AS t(i), "
+                                          "unnest(['Product description ' || i::VARCHAR]) AS products(description);");
+
+    ASSERT_FALSE(results->HasError()) << "Query failed: " << results->GetError();
+    ASSERT_EQ(results->RowCount(), 1);
+    nlohmann::json parsed = nlohmann::json::parse(results->GetValue(0, 0).GetValue<std::string>());
+    EXPECT_EQ(parsed.size(), 1);
+    EXPECT_EQ(parsed[0]["data"].size(), 1);
+}
+
 // Test GROUP BY with multiple tuples per group: LLM is called for each group
 TEST_F(LLMFirstTest, GroupByWithMultipleTuplesPerGroup) {
     EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, ::testing::_, ::testing::_, ::testing::_))

test/unit/functions/aggregate/llm_reduce.cpp

@@ -83,6 +83,31 @@ TEST_F(LLMReduceTest, MultipleTuplesWithoutGroupBy) {
     ASSERT_EQ(results->GetValue(0, 0).GetValue<std::string>(), GetExpectedResponse());
 }
 
+TEST_F(LLMReduceTest, DefaultBatchSizeSplitsLargeInput) {
+    constexpr size_t input_count = DEFAULT_BATCH_SIZE + 1;
+    const nlohmann::json first_batch_response = {{"items", {"Partial summary"}}};
+
+    EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, ::testing::_, ::testing::_, ::testing::_))
+            .Times(2);
+    EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+            .WillOnce(::testing::Return(std::vector<nlohmann::json>{first_batch_response}))
+            .WillOnce(::testing::Return(std::vector<nlohmann::json>{GetExpectedJsonResponse()}));
+
+    auto con = Config::GetConnection();
+
+    const auto results = con.Query(
+            "SELECT llm_reduce("
+            "{'model_name': 'gpt-4o'}, "
+            "{'prompt': 'Summarize the following product descriptions', 'context_columns': [{'data': description}]}"
+            ") AS product_summary FROM range(" +
+            std::to_string(input_count) + ") AS t(i), "
+                                          "unnest(['Product description ' || i::VARCHAR]) AS products(description);");
+
+    ASSERT_FALSE(results->HasError()) << "Query failed: " << results->GetError();
+    ASSERT_EQ(results->RowCount(), 1);
+    ASSERT_EQ(results->GetValue(0, 0).GetValue<std::string>(), GetExpectedResponse());
+}
+
 // Test GROUP BY with multiple tuples per group: LLM is called for each group
 TEST_F(LLMReduceTest, GroupByWithMultipleTuplesPerGroup) {
     EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, ::testing::_, ::testing::_, ::testing::_))

test/unit/functions/aggregate/llm_rerank.cpp

@@ -37,6 +37,12 @@ class LLMRerankTest : public LLMAggregateTestBase<LlmRerank> {
         return nlohmann::json{{"items", ranking_indices}};
     }
 
+    nlohmann::json PrepareSequentialRanking(size_t input_count) const {
+        std::vector<int> ranking_indices(input_count);
+        std::iota(ranking_indices.begin(), ranking_indices.end(), 0);
+        return nlohmann::json{{"items", ranking_indices}};
+    }
+
     std::string FormatExpectedResult(const nlohmann::json& response) const override {
         return response.dump();
     }
@@ -91,6 +97,38 @@ TEST_F(LLMRerankTest, MultipleTuplesWithoutGroupBy) {
     });
 }
 
+TEST_F(LLMRerankTest, DefaultBatchSizeSplitsLargeInput) {
+    constexpr size_t input_count = DEFAULT_BATCH_SIZE + 1;
+
+    {
+        ::testing::InSequence sequence;
+        EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, DEFAULT_BATCH_SIZE, ::testing::_, ::testing::_))
+                .Times(1);
+        EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+                .WillOnce(::testing::Return(std::vector<nlohmann::json>{PrepareSequentialRanking(DEFAULT_BATCH_SIZE)}));
+        EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, 9, ::testing::_, ::testing::_))
+                .Times(1);
+        EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+                .WillOnce(::testing::Return(std::vector<nlohmann::json>{PrepareSequentialRanking(9)}));
+    }
+
+    auto con = Config::GetConnection();
+
+    const auto results = con.Query(
+            "SELECT llm_rerank("
+            "{'model_name': 'gpt-4o'}, "
+            "{'prompt': 'Rank these products by relevance', 'context_columns': [{'data': description}]}"
+            ") AS reranked_products FROM range(" +
+            std::to_string(input_count) + ") AS t(i), "
+                                          "unnest(['Product description ' || i::VARCHAR]) AS products(description);");
+
+    ASSERT_FALSE(results->HasError()) << "Query failed: " << results->GetError();
+    ASSERT_EQ(results->RowCount(), 1);
+    nlohmann::json parsed = nlohmann::json::parse(results->GetValue(0, 0).GetValue<std::string>());
+    ASSERT_EQ(parsed.size(), 1);
+    EXPECT_EQ(parsed[0]["data"].size(), input_count);
+}
+
 // Test GROUP BY with multiple tuples per group: LLM is called for each group
 TEST_F(LLMRerankTest, GroupByWithMultipleTuplesPerGroup) {
     nlohmann::json response_2_items = nlohmann::json{{"items", {1, 0}}};

test/unit/functions/scalar/llm_complete.cpp

@@ -31,6 +31,14 @@ class LLMCompleteTest : public LLMFunctionTestBase<LlmComplete> {
         return expected_response;
     }
 
+    nlohmann::json PrepareExpectedResponseRange(size_t start_index, size_t count) const {
+        nlohmann::json expected_response = {{"items", {}}};
+        for (size_t i = 0; i < count; i++) {
+            expected_response["items"].push_back("response " + std::to_string(start_index + i));
+        }
+        return expected_response;
+    }
+
     std::string FormatExpectedResult(const nlohmann::json& response) const override {
         if (response.contains("items") && response["items"].is_array() && !response["items"].empty()) {
             return response["items"][0].get<std::string>();
@@ -132,11 +140,20 @@ TEST_F(LLMCompleteTest, Operation_LargeInputSet_ProcessesCorrectly) {
     constexpr size_t input_count = 100;
 
     const nlohmann::json expected_response = PrepareExpectedResponseForLargeInput(input_count);
+    std::vector<nlohmann::json> batch_responses;
+    for (size_t start_index = 0; start_index < input_count; start_index += DEFAULT_BATCH_SIZE) {
+        const auto batch_count = std::min<size_t>(DEFAULT_BATCH_SIZE, input_count - start_index);
+        batch_responses.push_back(PrepareExpectedResponseRange(start_index, batch_count));
+    }
+    size_t next_response = 0;
 
     EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, ::testing::_, ::testing::_, ::testing::_))
-            .Times(1);
+            .Times(static_cast<int>(batch_responses.size()));
     EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
-            .WillOnce(::testing::Return(std::vector<nlohmann::json>{expected_response}));
+            .Times(static_cast<int>(batch_responses.size()))
+            .WillRepeatedly(::testing::Invoke([&batch_responses, &next_response](const std::string&) {
+                return std::vector<nlohmann::json>{batch_responses[next_response++]};
+            }));
 
     // Use SQL approach - DuckDB's best practice for large datasets
     auto con = Config::GetConnection();
@@ -262,4 +279,4 @@ TEST_F(LLMCompleteTest, LLMCompleteAudioMissingTranscriptionModel) {
     ASSERT_TRUE(results->HasError());
 }
 
-}// namespace flock
+}// namespace flock

test/unit/functions/scalar/llm_embedding.cpp

@@ -154,4 +154,40 @@ TEST_F(LLMEmbeddingTest, Operation_LargeInputSet_ProcessesCorrectly) {
     }
 }
 
+TEST_F(LLMEmbeddingTest, Operation_DefaultBatchSizeSplitsLargeInput) {
+    constexpr size_t input_count = DEFAULT_BATCH_SIZE + 1;
+
+    nlohmann::json expected_response = nlohmann::json::array();
+    for (size_t i = 0; i < input_count; i++) {
+        std::vector<double> embedding;
+        for (size_t j = 0; j < 5; j++) {
+            embedding.push_back(0.01 * i + 0.1 * j);
+        }
+        expected_response.push_back(embedding);
+    }
+
+    {
+        ::testing::InSequence sequence;
+        EXPECT_CALL(*mock_provider, AddEmbeddingRequest(::testing::SizeIs(DEFAULT_BATCH_SIZE)))
+                .Times(1);
+        EXPECT_CALL(*mock_provider, AddEmbeddingRequest(::testing::SizeIs(1)))
+                .Times(1);
+    }
+    EXPECT_CALL(*mock_provider, CollectEmbeddings(::testing::_))
+            .WillOnce(::testing::Return(std::vector<nlohmann::json>{expected_response}));
+
+    auto con = Config::GetConnection();
+    const auto results = con.Query(
+            "SELECT " + GetFunctionName() + "("
+                                        "{'model_name': 'text-embedding-3-small'}, "
+                                        "{'context_columns': [{'data': content}]}"
+                                        ") AS embedding FROM range(" +
+            std::to_string(input_count) + ") AS t(i), unnest(['Document content number ' || i::VARCHAR]) AS tbl(content);");
+
+    ASSERT_TRUE(!results->HasError()) << "Query failed: " << results->GetError();
+    ASSERT_EQ(results->RowCount(), input_count);
+    ASSERT_EQ(results->GetValue(0, 0).type().id(), duckdb::LogicalTypeId::LIST);
+    ASSERT_EQ(results->GetValue(0, DEFAULT_BATCH_SIZE).type().id(), duckdb::LogicalTypeId::LIST);
+}
+
 }// namespace flock

test/unit/functions/scalar/llm_filter.cpp

@@ -130,6 +130,41 @@ TEST_F(LLMFilterTest, Operation_LargeInputSet_ProcessesCorrectly) {
     }
 }
 
+TEST_F(LLMFilterTest, Operation_DefaultBatchSizeSplitsLargeInput) {
+    constexpr size_t input_count = DEFAULT_BATCH_SIZE + 1;
+
+    nlohmann::json first_batch_response = {{"items", {}}};
+    for (size_t i = 0; i < DEFAULT_BATCH_SIZE; i++) {
+        first_batch_response["items"].push_back(i % 2 == 0);
+    }
+    nlohmann::json second_batch_response = {{"items", {false}}};
+
+    {
+        ::testing::InSequence sequence;
+        EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, DEFAULT_BATCH_SIZE, ::testing::_, ::testing::_))
+                .Times(1);
+        EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+                .WillOnce(::testing::Return(std::vector<nlohmann::json>{first_batch_response}));
+        EXPECT_CALL(*mock_provider, AddCompletionRequest(::testing::_, 1, ::testing::_, ::testing::_))
+                .Times(1);
+        EXPECT_CALL(*mock_provider, CollectCompletions(::testing::_))
+                .WillOnce(::testing::Return(std::vector<nlohmann::json>{second_batch_response}));
+    }
+
+    auto con = Config::GetConnection();
+    const auto results = con.Query(
+            "SELECT " + GetFunctionName() + "("
+                                        "{'model_name': 'gpt-4o'}, "
+                                        "{'prompt': 'Is this content relevant?', 'context_columns': [{'data': content}]}"
+                                        ") AS result FROM range(" +
+            std::to_string(input_count) + ") AS t(i), unnest(['Content item ' || i::VARCHAR]) AS tbl(content);");
+
+    ASSERT_TRUE(!results->HasError()) << "Query failed: " << results->GetError();
+    ASSERT_EQ(results->RowCount(), input_count);
+    EXPECT_EQ(results->GetValue(0, 0).GetValue<std::string>(), "true");
+    EXPECT_EQ(results->GetValue(0, DEFAULT_BATCH_SIZE).GetValue<std::string>(), "false");
+}
+
 // Test llm_filter with audio transcription
 TEST_F(LLMFilterTest, LLMFilterWithAudioTranscription) {
     const nlohmann::json expected_transcription = "{\"text\": \"This audio contains positive sentiment\"}";

test/unit/model_manager/model_manager_test.cpp

@@ -63,6 +63,21 @@ TEST_F(ModelManagerTest, ModelInitializationMinimal) {
         EXPECT_EQ(details.model_name, "gpt-4o-test");
         EXPECT_EQ(details.model, "gpt-4o");
         EXPECT_EQ(details.provider_name, "openai");
+        EXPECT_EQ(details.batch_size, 32);
+    });
+}
+
+TEST_F(ModelManagerTest, ModelInitializationUsesDefaultBatchSizeWhenUnset) {
+    json model_config = {
+            {"model_name", "gpt-4o"}};
+
+    EXPECT_NO_THROW({
+        Model model(model_config);
+        ModelDetails details = model.GetModelDetails();
+        EXPECT_EQ(details.model_name, "gpt-4o");
+        EXPECT_EQ(details.model, "gpt-4o");
+        EXPECT_EQ(details.provider_name, "openai");
+        EXPECT_EQ(details.batch_size, DEFAULT_BATCH_SIZE);
     });
 }
 
@@ -127,4 +142,4 @@ TEST_F(ModelManagerTest, GetModelDetails) {
     EXPECT_EQ(details.batch_size, 10);
 }
 
-}// namespace flock
+}// namespace flock