Included a thirdparty library with LLM predictor

Author: ThatUdeshUK

CMakeLists.txt

@@ -373,7 +373,7 @@ option(BUILD_BENCHMARKS "Enable building of the benchmark suite." FALSE)
 option(BUILD_TPCE "Enable building of the TPC-E tool." FALSE)
 option(DISABLE_BUILTIN_EXTENSIONS "Disable linking extensions." FALSE)
 option(ENABLE_PREDICT "Build the ML Operator with DuckDB" FALSE)
-option(USE_TORCH "Build the ML Operator with DuckDB using the libtorch predictor" FALSE)
+option(PREDICTOR_IMPL "Build the ML Operator with DuckDB using specified predictor (torchscript, onnx, llm_api)" "")
 option(BUILD_PYTHON "Build the DuckDB Python extension" FALSE)
 option(USER_SPACE "Build the DuckDB Python in the user space" FALSE)
 option(FORCE_QUERY_LOG "If enabled, all queries will be logged to the specified path" OFF)
@@ -1214,19 +1214,24 @@ endif()
 
 if (ENABLE_PREDICT)
   add_definitions(-DENABLE_PREDICT)
-  if (USE_TORCH)
+  if (PREDICTOR_IMPL STREQUAL "torchscript")
     if (NOT DEFINED ENV{TORCH_INSTALL_PREFIX})
       message( FATAL_ERROR "TORCH_INSTALL_PREFIX environment variable is not defined with USE_TORCH=1 option, CMake will exit." )
     endif()
     set(LIBTORCH_PATH "$ENV{TORCH_INSTALL_PREFIX}")
     set(LIBTORCH_INCLUDE_DIRS "${LIBTORCH_PATH}/include")
     message(STATUS "Torch Include Libs: ${LIBTORCH_INCLUDE_DIRS}")
-    add_definitions(-DUSE_TORCH)
-  else ()
+    add_definitions(-DPREDICTOR_IMPL=1)
+  endif ()
+  if (PREDICTOR_IMPL STREQUAL "onnx")
     if (NOT DEFINED ENV{ONNX_INSTALL_PREFIX})
       message( FATAL_ERROR "ONNX_INSTALL_PREFIX environment variable is not defined with USE_TORCH=0 option, CMake will exit." )
     endif()
     find_path(ONNX_RUNTIME_SESSION_INCLUDE_DIRS onnxruntime_cxx_api.h HINTS "$ENV{ONNX_INSTALL_PREFIX}/include")
+      add_definitions(-DPREDICTOR_IMPL=2)
+  endif()
+  if (PREDICTOR_IMPL STREQUAL "llm_api")
+    add_definitions(-DPREDICTOR_IMPL=3)
   endif()
 endif()
 

Makefile

@@ -147,10 +147,7 @@ ifeq (${BUILD_PYTHON}, 1)
 	CMAKE_VARS:=${CMAKE_VARS} -DBUILD_PYTHON=1 -DDUCKDB_EXTENSION_CONFIGS="tools/pythonpkg/duckdb_extension_config.cmake"
 endif
 ifeq (${ENABLE_PREDICT}, 1)
-	CMAKE_VARS:=${CMAKE_VARS} -DENABLE_PREDICT=1
-endif
-ifeq (${USE_TORCH}, 1)
-	CMAKE_VARS:=${CMAKE_VARS} -DUSE_TORCH=1
+	CMAKE_VARS:=${CMAKE_VARS} -DENABLE_PREDICT=1  -DPREDICTOR_IMPL="${PREDICTOR_IMPL}"
 endif
 ifeq (${PYTHON_USER_SPACE}, 1)
 	CMAKE_VARS:=${CMAKE_VARS} -DUSER_SPACE=1

src/CMakeLists.txt

@@ -111,10 +111,12 @@ else()
       duckdb_yyjson)
 
   if(ENABLE_PREDICT)
-    if(USE_TORCH)
+    if (${PREDICTOR_IMPL} STREQUAL "torchscript")
       set(DUCKDB_LINK_LIBS ${DUCKDB_LINK_LIBS} duckdb_torch)
-    else()
+    elseif (${PREDICTOR_IMPL} STREQUAL "onnx")
       set(DUCKDB_LINK_LIBS ${DUCKDB_LINK_LIBS} duckdb_onnx)
+    elseif (${PREDICTOR_IMPL} STREQUAL "llm_api")
+      set(DUCKDB_LINK_LIBS ${DUCKDB_LINK_LIBS} duckdb_llm_api)
     endif()
   endif()
 

src/common/enum_util.cpp

@@ -4614,6 +4614,8 @@ const char* EnumUtil::ToChars<ModelType>(ModelType value) {
 		return "LLM";
 	case ModelType::GNN:
 		return "GNN";
+	case ModelType::LLM_API:
+		return "LLM_API";
 	default:
 		throw NotImplementedException(StringUtil::Format("Enum value: '%d' not implemented", value));
 	}
@@ -4630,6 +4632,9 @@ ModelType EnumUtil::FromString<ModelType>(const char *value) {
 	if (StringUtil::Equals(value, "GNN")) {
 		return ModelType::GNN;
 	}
+	if (StringUtil::Equals(value, "LLM_API")) {
+		return ModelType::LLM_API;
+	}
 	throw NotImplementedException(StringUtil::Format("Enum value: '%s' not implemented", value));
 }
 

src/execution/CMakeLists.txt

@@ -1,13 +1,21 @@
 if(ENABLE_PREDICT)
-  if(USE_TORCH)
+  if (PREDICTOR_IMPL STREQUAL "torchscript")
+    message(
+      STATUS "Added TorchScript include directories")
     include_directories(${LIBTORCH_INCLUDE_DIRS})
     include_directories(../../third_party/predictors/torchscript/include)
-  else()
+  endif ()
+  if (PREDICTOR_IMPL STREQUAL "onnx")
     message(
-      STATUS "ORT include lib in src -> ${ONNX_RUNTIME_SESSION_INCLUDE_DIRS}")
+      STATUS "Added ORT include lib in src -> ${ONNX_RUNTIME_SESSION_INCLUDE_DIRS}")
     include_directories(${ONNX_RUNTIME_SESSION_INCLUDE_DIRS})
     include_directories(../../third_party/predictors/onnx/include)
   endif()
+  if (PREDICTOR_IMPL STREQUAL "llm_api")
+    message(
+      STATUS "Added LLM API include directories")
+    include_directories(../../third_party/predictors/llm_api/include)
+  endif()
 endif()
 
 add_subdirectory(expression_executor)

src/execution/operator/projection/physical_gnn_predict.cpp

@@ -9,7 +9,7 @@
 #include <iostream>
 #include <map>
 
-#if ENABLE_PREDICT
+#if defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 2
 #include "duckdb_onnx.hpp"
 #endif
 
@@ -67,7 +67,7 @@ PhysicalGNNPredict::PhysicalGNNPredict(vector<LogicalType> types_p, idx_t estima
 }
 
 unique_ptr<Predictor> PhysicalGNNPredict::InitPredictor() const {
-#if defined(ENABLE_PREDICT)
+#if defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 2
 	return make_uniq<ONNXPredictor>();
 #else
 	return nullptr;

src/execution/operator/projection/physical_predict.cpp

@@ -3,10 +3,12 @@
 #include <iostream>
 #include <map>
 
-#if defined(ENABLE_PREDICT) && defined(USE_TORCH)
+#if defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 1
 #include "duckdb_torch.hpp"
-#elif ENABLE_PREDICT
+#elif defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 2
 #include "duckdb_onnx.hpp"
+#elif defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 3
+#include "duckdb_llm_api.hpp"
 #endif
 
 #define CHUNK_PRED 1
@@ -45,10 +47,12 @@ PhysicalPredict::PhysicalPredict(vector<LogicalType> types_p, unique_ptr<Physica
 }
 
 unique_ptr<Predictor> PhysicalPredict::InitPredictor() const {
-#if defined(ENABLE_PREDICT) && USE_TORCH
+#if defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 1
 	return make_uniq<TorchPredictor>();
-#elif defined(ENABLE_PREDICT)
+#elif defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 2
 	return make_uniq<ONNXPredictor>();
+#elif defined(ENABLE_PREDICT) && PREDICTOR_IMPL == 3
+	return make_uniq<LlmApiPredictor>();
 #else
 	return nullptr;
 #endif
@@ -83,6 +87,9 @@ OperatorResultType PhysicalPredict::Execute(ExecutionContext &context, DataChunk
 	} else if (predictor.task == PREDICT_LLM_TASK) {
 		predictor.PredictLMChunk(input, predictions, (int)input.size(), this->input_mask, (int)result_set_types.size(),
 		                         state.stats);
+	} else if (predictor.task == PREDICT_LLM_API_TASK) {
+		predictor.PredictChunk(input, predictions, (int)input.size(), this->input_mask, (int)result_set_types.size(),
+		                       state.stats);
 	}
 #elif VEC_PRED
 	std::vector<float> inputs;

src/execution/physical_plan/plan_predict.cpp

@@ -9,7 +9,8 @@ namespace duckdb {
 unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalPredict &op) {
 	switch (op.bound_predict.model_type) {
 	case ModelType::TABULAR:
-	case ModelType::LLM: {
+	case ModelType::LLM:
+	case ModelType::LLM_API: {
 		D_ASSERT(op.children.size() == 1);
 		auto child_plan = CreatePlan(*op.children[0]);
 		auto predict = make_uniq<PhysicalPredict>(std::move(op.types), std::move(child_plan));

src/include/duckdb/common/enums/model_type.hpp

@@ -17,7 +17,8 @@ namespace duckdb {
 enum class ModelType : uint8_t {
     TABULAR,
     LLM,
-    GNN
+    GNN,
+    LLM_API,
 };
 
 } // namespace duckdb

src/include/duckdb/execution/operator/projection/physical_predict.hpp

@@ -15,7 +15,12 @@
 #include "duckdb/planner/expression.hpp"
 
 namespace duckdb {
-typedef enum PredictorTask { PREDICT_TABULAR_TASK = 0, PREDICT_LLM_TASK = 1, PREDICT_GNN_TASK = 2 } PredictorTask;
+typedef enum PredictorTask { 
+	PREDICT_TABULAR_TASK = 0, 
+	PREDICT_LLM_TASK = 1, 
+	PREDICT_GNN_TASK = 2, 
+	PREDICT_LLM_API_TASK = 3 
+} PredictorTask;
 
 struct PredictStats {
 	long load;