change the aggregate functions' interface for prompt details

dsg-polymtl · Nov 16, 2024 · 227ce6b · 227ce6b
1 parent a04a76e
commit 227ce6b
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 18 deletions.
diff --git a/src/core/functions/aggregate/llm_agg.cpp b/src/core/functions/aggregate/llm_agg.cpp
@@ -96,15 +96,18 @@ void LlmAggOperation::Initialize(const AggregateFunction &, data_ptr_t state_p)
 
 void LlmAggOperation::Operation(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, Vector &states,
                                 idx_t count) {
-    search_query = inputs[0].GetValue(0).ToString();
-
-    if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+    if (inputs[0].GetType().id() != LogicalTypeId::STRUCT) {
         throw std::runtime_error("Expected a struct type for model details");
     }
-
-    auto model_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+    auto model_details_json = CastVectorOfStructsToJson(inputs[0], 1)[0];
     LlmAggOperation::model_details = ModelManager::CreateModelDetails(CoreModule::GetConnection(), model_details_json);
 
+    if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+        throw std::runtime_error("Expected a struct type for prompt details");
+    }
+    auto prompt_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+    search_query = CreatePromptDetails(CoreModule::GetConnection(), prompt_details_json).prompt;
+
     if (inputs[2].GetType().id() != LogicalTypeId::STRUCT) {
         throw std::runtime_error("Expected a struct type for prompt inputs");
     }
@@ -174,10 +177,21 @@ void LlmAggOperation::FirstOrLastFinalize<FirstOrLast::FIRST>(Vector &states, Ag
 
 void LlmAggOperation::SimpleUpdate(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count,
                                    data_ptr_t state_p, idx_t count) {
-    search_query = inputs[0].GetValue(0).ToString();
-    auto model_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+    if (inputs[0].GetType().id() != LogicalTypeId::STRUCT) {
+        throw std::runtime_error("Expected a struct type for model details");
+    }
+    auto model_details_json = CastVectorOfStructsToJson(inputs[0], 1)[0];
     LlmAggOperation::model_details = ModelManager::CreateModelDetails(CoreModule::GetConnection(), model_details_json);
 
+    if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+        throw std::runtime_error("Expected a struct type for prompt details");
+    }
+    auto prompt_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+    search_query = CreatePromptDetails(CoreModule::GetConnection(), prompt_details_json).prompt;
+
+    if (inputs[2].GetType().id() != LogicalTypeId::STRUCT) {
+        throw std::runtime_error("Expected a struct type for prompt inputs");
+    }
     auto tuples = CastVectorOfStructsToJson(inputs[2], count);
 
     auto state_map_p = reinterpret_cast<LlmAggState *>(state_p);

diff --git a/src/core/functions/aggregate/llm_first.cpp b/src/core/functions/aggregate/llm_first.cpp
@@ -9,7 +9,7 @@ namespace core {
 
 void CoreAggregateFunctions::RegisterLlmFirstFunction(DatabaseInstance &db) {
     auto string_concat =
-        AggregateFunction("llm_first", {LogicalType::VARCHAR, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
+        AggregateFunction("llm_first", {LogicalType::ANY, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
                           AggregateFunction::StateSize<LlmAggState>, LlmAggOperation::Initialize,
                           LlmAggOperation::Operation, LlmAggOperation::Combine,
                           LlmAggOperation::FirstOrLastFinalize<FirstOrLast::FIRST>, LlmAggOperation::SimpleUpdate);

diff --git a/src/core/functions/aggregate/llm_last.cpp b/src/core/functions/aggregate/llm_last.cpp
@@ -9,7 +9,7 @@ namespace core {
 
 void CoreAggregateFunctions::RegisterLlmLastFunction(DatabaseInstance &db) {
     auto string_concat =
-        AggregateFunction("llm_last", {LogicalType::VARCHAR, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
+        AggregateFunction("llm_last", {LogicalType::ANY, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
                           AggregateFunction::StateSize<LlmAggState>, LlmAggOperation::Initialize,
                           LlmAggOperation::Operation, LlmAggOperation::Combine,
                           LlmAggOperation::FirstOrLastFinalize<FirstOrLast::LAST>, LlmAggOperation::SimpleUpdate);

diff --git a/src/core/functions/aggregate/llm_reduce.cpp b/src/core/functions/aggregate/llm_reduce.cpp
@@ -96,16 +96,19 @@ struct LlmReduceOperation {
 
     static void Operation(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, Vector &states,
                           idx_t count) {
-        reduce_query = inputs[0].GetValue(0).ToString();
-
-        if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+        if (inputs[0].GetType().id() != LogicalTypeId::STRUCT) {
             throw std::runtime_error("Expected a struct type for model details");
         }
-
-        auto model_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+        auto model_details_json = CastVectorOfStructsToJson(inputs[0], 1)[0];
         LlmReduceOperation::model_details =
             ModelManager::CreateModelDetails(CoreModule::GetConnection(), model_details_json);
 
+        if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+            throw std::runtime_error("Expected a struct type for prompt details");
+        }
+        auto prompt_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+        reduce_query = CreatePromptDetails(CoreModule::GetConnection(), prompt_details_json).prompt;
+
         if (inputs[2].GetType().id() != LogicalTypeId::STRUCT) {
             throw std::runtime_error("Expected a struct type for prompt inputs");
         }
@@ -160,11 +163,22 @@ struct LlmReduceOperation {
 
     static void SimpleUpdate(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count,
                              data_ptr_t state_p, idx_t count) {
-        reduce_query = inputs[0].GetValue(0).ToString();
-        auto model_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+        if (inputs[0].GetType().id() != LogicalTypeId::STRUCT) {
+            throw std::runtime_error("Expected a struct type for model details");
+        }
+        auto model_details_json = CastVectorOfStructsToJson(inputs[0], 1)[0];
         LlmReduceOperation::model_details =
             ModelManager::CreateModelDetails(CoreModule::GetConnection(), model_details_json);
 
+        if (inputs[1].GetType().id() != LogicalTypeId::STRUCT) {
+            throw std::runtime_error("Expected a struct type for prompt details");
+        }
+        auto prompt_details_json = CastVectorOfStructsToJson(inputs[1], 1)[0];
+        reduce_query = CreatePromptDetails(CoreModule::GetConnection(), prompt_details_json).prompt;
+
+        if (inputs[2].GetType().id() != LogicalTypeId::STRUCT) {
+            throw std::runtime_error("Expected a struct type for prompt inputs");
+        }
         auto tuples = CastVectorOfStructsToJson(inputs[2], count);
 
         auto state_map_p = reinterpret_cast<LlmAggState *>(state_p);
@@ -187,7 +201,7 @@ std::unordered_map<void *, std::shared_ptr<LlmAggState>> LlmReduceOperation::sta
 
 void CoreAggregateFunctions::RegisterLlmReduceFunction(DatabaseInstance &db) {
     auto string_concat = AggregateFunction(
-        "llm_reduce", {LogicalType::VARCHAR, LogicalType::ANY, LogicalType::ANY}, LogicalType::VARCHAR,
+        "llm_reduce", {LogicalType::ANY, LogicalType::ANY, LogicalType::ANY}, LogicalType::VARCHAR,
         AggregateFunction::StateSize<LlmAggState>, LlmReduceOperation::Initialize, LlmReduceOperation::Operation,
         LlmReduceOperation::Combine, LlmReduceOperation::Finalize, LlmReduceOperation::SimpleUpdate);
 

diff --git a/src/core/functions/aggregate/llm_rerank.cpp b/src/core/functions/aggregate/llm_rerank.cpp
@@ -123,7 +123,7 @@ void LlmAggOperation::RerankerFinalize(Vector &states, AggregateInputData &aggr_
 
 void CoreAggregateFunctions::RegisterLlmRerankFunction(DatabaseInstance &db) {
     auto string_concat = AggregateFunction(
-        "llm_rerank", {LogicalType::VARCHAR, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
+        "llm_rerank", {LogicalType::ANY, LogicalType::ANY, LogicalType::ANY}, LogicalType::JSON(),
         AggregateFunction::StateSize<LlmAggState>, LlmAggOperation::Initialize, LlmAggOperation::Operation,
         LlmAggOperation::Combine, LlmAggOperation::RerankerFinalize, LlmAggOperation::SimpleUpdate);