1 line
20 KiB
JSON
1 line
20 KiB
JSON
{"repo": "asu-cactus/cactusdb", "files": {"velox/ml_functions/ChatGPT.h": {"4fb6c8afc9c830e4975194ea321f91c8956943fd": {"t": 1785595759.3878238, "n": 19631, "h": "55e2f5cd07b0f788", "c": "/*\n * Copyright (c) 2025 ASU Cactus Lab.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#pragma once\n#include <cpr/cpr.h>\n#include <fmt/format.h>\n#include <nlohmann/json.hpp>\n#include <chrono>\n#include <cstdint>\n#include <ctime>\n#include <fstream>\n#include <iomanip>\n#include <iostream>\n#include <thread>\n#include \"BaseFunction.h\"\n#include \"velox/exec/tests/utils/AssertQueryBuilder.h\"\n#include \"velox/exec/tests/utils/PlanBuilder.h\"\n#include \"velox/exec/tests/utils/TempDirectoryPath.h\"\n#include \"velox/expression/VectorFunction.h\"\n#include \"velox/ml_functions/UtilFunction.h\"\n#include \"velox/vector/tests/utils/VectorTestBase.h\"\n\nusing namespace facebook::velox;\nusing namespace facebook::velox::test;\nusing namespace facebook::velox::exec;\nusing namespace facebook::velox::exec::test;\nusing namespace facebook::velox::memory;\n\n#define MAX_ALLOWED_CHATGPT_TRY 30\n\n// Function to count the number of punctuation marks in a string\nint countPunctuation(const std::string& str) {\n int punctuationCount = 0;\n for (char ch : str) {\n if (std::ispunct(ch)) {\n ++punctuationCount;\n }\n }\n return punctuationCount;\n}\n\nvoid sendRequestViaCpr(\n const std::string& url,\n const cpr::Header& headers,\n const std::string& payload,\n cpr::Response& response,\n int& numFailures) {\n response = cpr::Post(cpr::Url{url}, cpr::Header{headers}, cpr::Body{payload});\n int failureCount = 0;\n while (response.status_code != 200) {\n response =\n cpr::Post(cpr::Url{url}, cpr::Header{headers}, cpr::Body{payload});\n if (++failureCount > MAX_ALLOWED_CHATGPT_TRY) {\n throw std::runtime_error(fmt::format(\n \"[ERROR] Failed to send request to OpenAI API. Reached maximum number of retries: {}, error message: {}\",\n MAX_ALLOWED_CHATGPT_TRY,\n payload));\n }\n }\n numFailures = failureCount;\n}\n\nvoid sendRequestViaCprBatch(\n const std::string& url,\n const cpr::Header& headers,\n const std::vector<std::string>& payloadVector,\n std::vector<cpr::Response>& responseVector,\n const int& startIdx,\n std::vector<int>& numFailureVector) {\n int index = startIdx;\n for (std::string payload : payloadVector) {\n int failureCount = 0;\n auto response =\n cpr::Post(cpr::Url{url}, cpr::Header{headers}, cpr::Body{payload});\n while (response.status_code != 200) {\n response =\n cpr::Post(cpr::Url{url}, cpr::Header{headers}, cpr::Body{payload});\n if (++failureCount > MAX_ALLOWED_CHATGPT_TRY) {\n throw std::runtime_error(fmt::format(\n \"[ERROR] Failed to send request to OpenAI API. Reached maximum number of retries: {}, error message: {}\",\n MAX_ALLOWED_CHATGPT_TRY,\n payload));\n }\n }\n responseVector[index] = response;\n numFailureVector[index] = failureCount;\n index++;\n }\n}\n\nclass ChatGPT : public MLFunction {\n public:\n ChatGPT() {\n apiKey_ = getEnvVar(\"OPENAI_API_KEY\");\n if (apiKey_ == \"\") {\n throw std::runtime_error(fmt::format(\n \"[ERROR] OpenAI API key is not set, please set OPENAI_API_KEY\"));\n }\n numThreads_ = getEnvVar(\"NUM_THREADS\") == \"\"\n ? 8\n : std::stoi(getEnvVar(\"NUM_THREADS\"));\n url_ = \"https://api.openai.com/v1/chat/completions\";\n model_ = \"gpt-3.5-turbo\";\n inputTokenNumber_ = 0;\n outputTokenNumber_ = 0;\n numFailures_ = 0;\n }\n\n ChatGPT(std::string url, std::string model) {\n apiKey_ = getEnvVar(\"OPENAI_API_KEY\");\n if (apiKey_ == \"\") {\n throw std::runtime_error(fmt::format(\n \"[ERROR] OpenAI API key is not set, please set OPENAI_API_KEY\"));\n }\n numThreads_ = getEnvVar(\"NUM_THREADS\") == \"\"\n ? 8\n : std::stoi(getEnvVar(\"NUM_THREADS\"));\n url_ = url;\n model_ = model;\n inputTokenNumber_ = 0;\n outputTokenNumber_ = 0;\n numFailures_ = 0;\n }\n\n ~ChatGPT() {\n std::string filename = \"chatgpt.log\";\n std::ofstream file(filename, std::ios::app);\n if (!file) {\n std::cerr << \"Unable to open file: \" << filename << std::endl;\n return;\n }\n // Get the current time\n auto now = std::chrono::system_clock::now();\n std::time_t now_c = std::chrono::system_clock::to_time_t(now);\n\n // Write the timestamp to the file\n file << std::put_time(std::localtime(&now_c), \"%Y-%m-%d %H:%M:%S\") << \" \";\n // Write the uint64_t values to the file\n file << \"[ChatGPT] # Input:\" << inputTokenNumber_\n << \" # Output: \" << outputTokenNumber_\n << \" # NumFailure: \" << numFailures_ << std::endl;\n file.close();\n }\n\n void apply(\n const SelectivityVector& rows,\n std::vector<VectorPtr>& args,\n const TypePtr& type,\n exec::EvalCtx& context,\n VectorPtr& output) const override {\n std::string promptPrefix = \"\";\n BaseVector::ensureWritable(rows, type, context.pool(), output);\n\n exec::LocalDecodedVector decodedStringHolder(context, *args[0], rows);\n auto decodedStringInput = decodedStringHolder.get();\n\n int numInput = rows.size();\n int numSelected = rows.countSelected();\n LOG(INFO) << \"[INFO ChatGPT:] countSelected: \" << rows.countSelected()\n << \" numInput: \" << numInput << std::endl;\n\n if (args.size() == 2) {\n exec::LocalDecodedVector decodedStringHolder2(context, *args[1], rows);\n auto decodedStringInput2 = decodedStringHolder2.get();\n StringView val = decodedStringInput2->valueAt<StringView>(0);\n promptPrefix = std::string(val);\n }\n\n std::vector<std::string> results;\n\n cpr::Header headers{\n {\"Content-Type\", \"application/json\"},\n {\"Authorization\", \"Bearer \" + apiKey_}};\n\n // Thread vector\n std::vector<std::thread> threads;\n std::vector<cpr::Response> responses(numInput);\n std::vector<int> numFailureVector(numInput);\n\n int numInputsPerThread = int(std::ceil(float(numSelected) / numThreads_));\n int processedInputCount = 0;\n std::vector<std::string> payloadsBatchVector;\n int processedIndex = 0;\n\n // Version 1\n // This approach is more efficient by sending requests in batches and\n // leveraging multiple threads to send requests concurrently, it requires\n // additional isValid check to skip the rows that are not selected. Note: at\n // the end of this approach, it is required to invoke\n // context.moveOrCopyResult to copy the results back to the output vector\n // since we only compute the results for selected ones\n for (int i = 0; i < numInput; i++) {\n // if the row is not selected, skip\n if (!rows.isValid(i)) {\n continue;\n }\n StringView val = decodedStringInput->valueAt<StringView>(i);\n std::string valString = promptPrefix + std::string(val);\n nlohmann::json messageArrays = nlohmann::json::array();\n // Add message\n messageArrays.push_back({{\"role\", \"user\"}, {\"content\", valString}});\n\n nlohmann::json payload = {\n {\"model\", model_}, {\"messages\", messageArrays}, {\"max_tokens\", 150}};\n\n payloadsBatchVector.push_back(payload.dump());\n processedInputCount++;\n\n if (processedInputCount == numInputsPerThread || i == numInput - 1) {\n threads.emplace_back(\n sendRequestViaCprBatch,\n url_,\n headers,\n payloadsBatchVector,\n std::ref(responses),\n processedIndex - processedInputCount + 1,\n std::ref(numFailureVector));\n processedInputCount = 0;\n payloadsBatchVector.clear();\n }\n processedIndex++;\n }\n\n for (auto& thread : threads) {\n thread.join();\n }\n\n for (int i = 0; i < numSelected; i++) {\n if (responses[i].status_code == 200) {\n // parse the returned value\n nlohmann::json response_json = nlohmann::json::parse(responses[i].text);\n std::string generated_message =\n response_json[\"choices\"][0][\"message\"][\"content\"];\n results.push_back(generated_message);\n const_cast<uint64_t&>(inputTokenNumber_) = inputTokenNumber_ +\n response_json[\"usage\"][\"prompt_tokens\"].get<int>();\n const_cast<uint64_t&>(outputTokenNumber_) = outputTokenNumber_ +\n response_json[\"usage\"][\"completion_tokens\"].get<int>();\n const_cast<uint64_t&>(numFailures_) =\n numFailures_ + numFailureVector[i];\n if (numFailureVector[i] > 0) {\n LOG(WARNING)\n << \"[WARNING] Failed to send request to OpenAI API. Number of retries: \"\n << numFailureVector[i] << \" numFailures_: \" << numFailures_\n << std::endl;\n }\n LOG(INFO) << fmt::format(\n \"[INFO] i: {} / {}, results: {}, numFailures: {}\",\n i + 1,\n numSelected,\n generated_message,\n numFailureVector[i])\n << std::endl;\n } else {\n LOG(ERROR) << \"Error: \" << responses[i].status_code << \" - \"\n << responses[i].text << std::endl;\n }\n }\n\n VectorMaker maker{context.pool()};\n VectorPtr localResult = maker.flatVector<std::string>(results);\n\n context.moveOrCopyResult(localResult, rows, output);\n\n // Version 2: Leveraging applyToSelected function, while this is done\n // sequentially, it is easier to implement and debug but less efficient\n /*\n auto flatResult = output->asFlatVector<StringView>();\n rows.applyToSelected([&](vector_size_t row) {\n StringView val = decodedStringInput->valueAt<StringView>(row);\n std::string valString = promptPrefix + std::string(val);\n nlohmann::json messageArrays = nlohmann::json::array();\n // Add message\n messageArrays.push_back({{\"role\", \"user\"}, {\"content\", valString}});\n\n nlohmann::json payload = {\n {\"model\", model_}, {\"messages\", messageArrays}, {\"max_tokens\", 150}};\n sendRequestViaCpr(url_, headers, payload.dump(), responses[row],\n numFailureVector[row]);\n\n // parse the returned value\n nlohmann::json response_json = nlohmann::json::parse(responses[row].text);\n std::string generated_message =\n response_json[\"choices\"][0][\"message\"][\"content\"];\n flatResult->set(row, StringView(generated_message));\n // results.push_back(generated_message);\n const_cast<uint64_t&>(inputTokenNumber_) = inputTokenNumber_ +\n response_json[\"usage\"][\"prompt_tokens\"].get<int>();\n const_cast<uint64_t&>(outputTokenNumber_) = outputTokenNumber_ +\n response_json[\"usage\"][\"completion_tokens\"].get<int>();\n const_cast<uint64_t&>(numFailures_) =\n numFailures_ + numFailureVector[row];\n LOG(INFO) << fmt::format(\n \"[INFO] Selected row: {} / {}, results: {},\n numFailures: {}\", row + 1, numSelected, generated_message,\n numFailureVector[row])\n << std::endl;\n });\n */\n }\n\n static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {\n return {\n exec::FunctionSignatureBuilder()\n .argumentType(\"VARCHAR\")\n .returnType(\"VARCHAR\")\n .build(),\n // supports with prompt prefix\n exec::FunctionSignatureBuilder()\n .argumentType(\"VARCHAR\")\n .argumentType(\"VARCHAR\")\n .returnType(\"VARCHAR\")\n .build()};\n }\n\n float* getTensor() const override {\n // TODO: need to implement\n return nullptr;\n }\n\n static std::string getName() {\n return \"chatgpt\";\n };\n\n CostEstimate getCost(std::vector<int> inputDims) {\n // TODO: need to implement\n return CostEstimate(0, inputDims[0], inputDims[1]);\n }\n\n private:\n std::string apiKey_;\n std::string url_;\n std::string model_;\n uint64_t inputTokenNumber_;\n uint64_t outputTokenNumber_;\n uint64_t numFailures_;\n int numThreads_;\n};\n\nclass ChatGPTRecommender : public MLFunction {\n public:\n ChatGPTRecommender() {\n apiKey_ = getEnvVar(\"OPENAI_API_KEY\");\n if (apiKey_ == \"\") {\n throw std::runtime_error(fmt::format(\n \"[ERROR] OpenAI API key is not set, please set OPENAI_API_KEY\"));\n }\n numThreads_ = getEnvVar(\"NUM_THREADS\") == \"\"\n ? 8\n : std::stoi(getEnvVar(\"NUM_THREADS\"));\n url_ = \"https://api.openai.com/v1/chat/completions\";\n model_ = \"gpt-3.5-turbo\";\n inputTokenNumber_ = 0;\n outputTokenNumber_ = 0;\n numFailures_ = 0;\n }\n\n ChatGPTRecommender(std::string url, std::string model) {\n apiKey_ = getEnvVar(\"OPENAI_API_KEY\");\n if (apiKey_ == \"\") {\n throw std::runtime_error(fmt::format(\n \"[ERROR] OpenAI API key is not set, please set OPENAI_API_KEY\"));\n }\n numThreads_ = getEnvVar(\"NUM_THREADS\") == \"\"\n ? 8\n : std::stoi(getEnvVar(\"NUM_THREADS\"));\n url_ = url;\n model_ = model;\n inputTokenNumber_ = 0;\n outputTokenNumber_ = 0;\n numFailures_ = 0;\n }\n\n ~ChatGPTRecommender() {\n std::string filename = \"chatgpt.log\";\n std::ofstream file(filename, std::ios::app);\n if (!file) {\n std::cerr << \"Unable to open file: \" << filename << std::endl;\n return;\n }\n // Get the current time\n auto now = std::chrono::system_clock::now();\n std::time_t now_c = std::chrono::system_clock::to_time_t(now);\n\n // Write the timestamp to the file\n file << std::put_time(std::localtime(&now_c), \"%Y-%m-%d %H:%M:%S\") << \" \";\n\n // Write the uint64_t values to the file\n file << \"[ChatGPT Recommender] # Input:\" << inputTokenNumber_\n << \" # Output: \" << outputTokenNumber_\n << \" # NumFailure: \" << numFailures_ << std::endl;\n file.close();\n }\n\n void apply(\n const SelectivityVector& rows,\n std::vector<VectorPtr>& args,\n const TypePtr& type,\n exec::EvalCtx& context,\n VectorPtr& output) const override {\n std::string promptSuffix = \"\";\n BaseVector::ensureWritable(rows, type, context.pool(), output);\n\n exec::LocalDecodedVector decodedStringHolder1(context, *args[0], rows);\n auto decodedStringInput1 = decodedStringHolder1.get();\n\n exec::LocalDecodedVector decodedStringHolder2(context, *args[1], rows);\n auto decodedStringInput2 = decodedStringHolder2.get();\n\n int numInput = rows.size();\n int numSelected = rows.countSelected();\n LOG(INFO) << \"[INFO ChatGPTRecommender:] countSelected: \"\n << rows.countSelected() << \" numInput: \" << numInput << std::endl;\n\n if (args.size() == 3) {\n exec::LocalDecodedVector decodedStringHolder3(context, *args[2], rows);\n auto decodedStringInput3 = decodedStringHolder3.get();\n StringView val = decodedStringInput3->valueAt<StringView>(0);\n promptSuffix = std::string(val);\n }\n\n std::vector<std::string> results;\n\n cpr::Header headers{\n {\"Content-Type\", \"application/json\"},\n {\"Authorization\", \"Bearer \" + apiKey_}};\n\n // Thread vector\n std::vector<std::thread> threads;\n std::vector<cpr::Response> responses(numInput);\n std::vector<int> numFailureVector(numInput);\n\n int numInputsPerThread = int(std::ceil(float(numSelected) / numThreads_));\n int processedInputCount = 0;\n std::vector<std::string> payloadsBatchVector;\n int processedIndex = 0;\n\n for (int i = 0; i < numInput; i++) {\n // if the row is not selected, skip\n if (!rows.isValid(i)) {\n continue;\n }\n StringView val1 = decodedStringInput1->valueAt<StringView>(i);\n StringView val2 = decodedStringInput2->valueAt<StringView>(i);\n std::string valString =\n \"Summarized user statistics data (preference): \" + std::string(val1) +\n \". \\n Summarized user movie metadata: \" + std::string(val2) + \".\\n\" +\n promptSuffix;\n nlohmann::json messageArrays = nlohmann::json::array();\n // Add message\n messageArrays.push_back({{\"role\", \"user\"}, {\"content\", valString}});\n\n nlohmann::json payload = {\n {\"model\", model_}, {\"messages\", messageArrays}, {\"max_tokens\", 500}};\n\n payloadsBatchVector.push_back(payload.dump());\n processedInputCount++;\n\n if (processedInputCount == numInputsPerThread || i == numInput - 1) {\n threads.emplace_back(\n sendRequestViaCprBatch,\n url_,\n headers,\n payloadsBatchVector,\n std::ref(responses),\n processedIndex - processedInputCount + 1,\n std::ref(numFailureVector));\n processedInputCount = 0;\n payloadsBatchVector.clear();\n }\n processedIndex++;\n }\n\n for (auto& thread : threads) {\n thread.join();\n }\n\n for (int i = 0; i < numSelected; i++) {\n if (responses[i].status_code == 200) {\n // parse the returned value\n nlohmann::json response_json = nlohmann::json::parse(responses[i].text);\n std::string generated_message =\n response_json[\"choices\"][0][\"message\"][\"content\"];\n results.push_back(generated_message);\n const_cast<uint64_t&>(inputTokenNumber_) = inputTokenNumber_ +\n response_json[\"usage\"][\"prompt_tokens\"].get<int>();\n const_cast<uint64_t&>(outputTokenNumber_) = outputTokenNumber_ +\n response_json[\"usage\"][\"completion_tokens\"].get<int>();\n const_cast<uint64_t&>(numFailures_) =\n numFailures_ + numFailureVector[i];\n if (numFailureVector[i] > 0) {\n LOG(WARNING)\n << \"[WARNING] Failed to send request to OpenAI API. Number of retries: \"\n << numFailureVector[i] << \" numFailures_: \" << numFailures_\n << std::endl;\n }\n LOG(INFO) << fmt::format(\n \"[INFO] i: {} / {}, results: {}, numFailures: {}\",\n i + 1,\n numSelected,\n generated_message,\n numFailureVector[i])\n << std::endl;\n } else {\n LOG(ERROR) << \"Error: \" << responses[i].status_code << \" - \"\n << responses[i].text << std::endl;\n }\n }\n\n VectorMaker maker{context.pool()};\n VectorPtr localResult = maker.flatVector<std::string>(results);\n\n context.moveOrCopyResult(localResult, rows, output);\n }\n\n static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {\n return {\n exec::FunctionSignatureBuilder()\n .argumentType(\"VARCHAR\")\n .argumentType(\"VARCHAR\")\n .returnType(\"VARCHAR\")\n .build(),\n // supports with prompt suffix\n exec::FunctionSignatureBuilder()\n .argumentType(\"VARCHAR\")\n .argumentType(\"VARCHAR\")\n .argumentType(\"VARCHAR\")\n .returnType(\"VARCHAR\")\n .build()};\n }\n\n float* getTensor() const override {\n // TODO: need to implement\n return nullptr;\n }\n\n static std::string getName() {\n return \"chatgpt_recommender\";\n };\n\n CostEstimate getCost(std::vector<int> inputDims) {\n // TODO: need to implement\n return CostEstimate(0, inputDims[0], inputDims[1]);\n }\n\n private:\n std::string apiKey_;\n std::string url_;\n std::string model_;\n uint64_t inputTokenNumber_;\n uint64_t outputTokenNumber_;\n uint64_t numFailures_;\n int numThreads_;\n};"}}}} |