299 lines
10 KiB
C++
299 lines
10 KiB
C++
/*
|
||
* Copyright (C) 2024–2025 Amlogic, Inc. All rights reserved.
|
||
*
|
||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||
* you may not use this file except in compliance with the License.
|
||
* You may obtain a copy of the License at
|
||
*
|
||
* http://www.apache.org/licenses/LICENSE-2.0
|
||
*
|
||
* Unless required by applicable law or agreed to in writing, software
|
||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
* See the License for the specific language governing permissions and
|
||
* limitations under the License.
|
||
*/
|
||
#include <cstdio>
|
||
#include <iostream>
|
||
#include <vector>
|
||
#include <string>
|
||
#include <cstring>
|
||
#include <regex>
|
||
|
||
#include "llmsdk.h"
|
||
#include <nlohmann/json.hpp>
|
||
|
||
// Small holder for callback behavior:
|
||
// - printed: print tokens as they come
|
||
// - buffer: collect first-pass output (for external two-run mode)
|
||
struct DemoUserData {
|
||
bool printed = true;
|
||
int request_id;
|
||
std::string buffer;
|
||
};
|
||
|
||
// Stream callback: either print tokens or accumulate to buffer
|
||
static void result_callback(AML_LLMResult* result, void* userdata, AML_LLMRunStatus run_status) {
|
||
|
||
// std::cout << result->text;
|
||
// std::cout.flush();
|
||
auto* mydata = reinterpret_cast<DemoUserData*>(userdata);
|
||
if (!mydata) return;
|
||
|
||
if (run_status == AML_LLM_RUN_NORMAL) {
|
||
if (result && result->text) {
|
||
if (mydata->printed) {
|
||
std::cout << result->text;
|
||
std::cout.flush();
|
||
}
|
||
mydata->buffer.append(result->text);
|
||
}
|
||
} else if (run_status == AML_LLM_RUN_FINISH) {
|
||
if (mydata->printed) std::cout << "\n[FINISHED]\n";
|
||
} else if (run_status == AML_LLM_RUN_ERROR) {
|
||
std::cerr << "\n[ERROR]\n";
|
||
}
|
||
}
|
||
|
||
// Internal tool execution callback (JNI-internal mode)
|
||
static int tool_callback(
|
||
const char* tool_name,
|
||
const char* arguments_json,
|
||
void* /*userdata*/,
|
||
char** out_result_buffer
|
||
) {
|
||
std::cout << "\n[TOOL CALL - internal] name=" << tool_name
|
||
<< " args=" << (arguments_json ? arguments_json : "(null)")
|
||
<< "\n";
|
||
|
||
if (strcmp(tool_name, "get_weather") == 0) {
|
||
const char* resp =
|
||
"{"
|
||
"\"city\":\"Hangzhou\","
|
||
"\"temperature\":\"25C\","
|
||
"\"condition\":\"Sunny\""
|
||
"}";
|
||
*out_result_buffer = strdup(resp); // Freed by SDK
|
||
return 0;
|
||
}
|
||
|
||
const char* fallback = "{\"error\":\"tool_not_found\"}";
|
||
*out_result_buffer = strdup(fallback);
|
||
return -1;
|
||
}
|
||
|
||
// External tool execution stub (for demo)
|
||
static std::string external_execute_tool(const std::string& tool_name, const std::string& /*args_json*/) {
|
||
if (tool_name == "get_weather") {
|
||
return std::string("{") +
|
||
"\"city\":\"Hangzhou\","
|
||
"\"temperature\":\"25C\","
|
||
"\"condition\":\"Sunny\""
|
||
"}";
|
||
}
|
||
return std::string("{\"error\":\"tool_not_found\"}");
|
||
}
|
||
|
||
// Parse <tool_call>{...}</tool_call> blocks from model output
|
||
struct ToolCall { std::string name; std::string args_json; };
|
||
static std::vector<ToolCall> parse_tool_calls(const std::string& text) {
|
||
std::vector<ToolCall> calls;
|
||
printf("model_output: %s\n", text.c_str());
|
||
|
||
// match across newlines safely
|
||
std::regex tool_call_pattern(
|
||
"<tool_call>[\\n\\r\\s]*(\\{[\\n\\r\\s\\S]*?\\})[\\n\\r\\s]*</tool_call>",
|
||
std::regex::icase | std::regex::ECMAScript
|
||
);
|
||
|
||
std::sregex_iterator it(text.begin(), text.end(), tool_call_pattern);
|
||
std::sregex_iterator end;
|
||
|
||
|
||
for (; it != end; ++it) {
|
||
std::string json_str = (*it)[1].str();
|
||
printf("Matched tool_call JSON: %s\n", json_str.c_str());
|
||
nlohmann::json parsed;
|
||
parsed = nlohmann::json::parse(json_str, nullptr, /* allow_exceptions = */ false);
|
||
if (parsed.is_discarded()) {
|
||
printf("Function Calling JSON parse failed: %s\n", json_str.c_str());
|
||
continue;
|
||
}
|
||
|
||
std::string tool;
|
||
if (parsed.contains("tool_name")) tool = parsed["tool_name"].get<std::string>();
|
||
else if (parsed.contains("name")) tool = parsed["name"].get<std::string>();
|
||
|
||
if (!tool.empty()) {
|
||
calls.push_back(ToolCall{tool, parsed["arguments"].dump()});
|
||
}
|
||
printf("tool: %s, arguments: %s\n", tool.c_str(), parsed["arguments"].dump().c_str());
|
||
|
||
}
|
||
return calls;
|
||
}
|
||
|
||
// Build JSON array of tool results by executing each tool externally
|
||
static std::string build_tool_results_json(const std::vector<ToolCall>& calls) {
|
||
std::string out = "[";
|
||
for (size_t i = 0; i < calls.size(); ++i) {
|
||
const auto& c = calls[i];
|
||
std::string one = external_execute_tool(c.name, c.args_json);
|
||
out += one;
|
||
if (i + 1 < calls.size()) out += ",";
|
||
}
|
||
out += "]";
|
||
return out;
|
||
}
|
||
|
||
int main(int argc, char** argv) {
|
||
// Init LLM context
|
||
LLMContext ctx = nullptr;
|
||
|
||
AML_LLMInitConfig init_cfg;
|
||
std::memset(&init_cfg, 0, sizeof(init_cfg));
|
||
init_cfg.model_path = argc > 1 ? argv[1] : "/path/to/your/model.bin"; // TODO: replace
|
||
init_cfg.tokenizer_path = argc > 2 ? argv[2] : "/path/to/your/tokenizer.model"; // TODO: replace
|
||
init_cfg.sampling_mode = AML_LLM_TOP_P;
|
||
init_cfg.top_k = 0;
|
||
init_cfg.top_p = 0.9f;
|
||
init_cfg.temperature = 0.7f;
|
||
init_cfg.repeat_penalty = 1.0f;
|
||
|
||
if (aml_llm_init(&ctx, &init_cfg, result_callback) != AML_LLM_Status_Success) {
|
||
std::cerr << "Init failed: aml_llm_init\n";
|
||
return -1;
|
||
}
|
||
|
||
// Define tools JSON exposed to model
|
||
const char* tools_json = R"JSON(
|
||
[
|
||
{
|
||
"name": "get_weather",
|
||
"description": "Get current weather for a city",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"city": {
|
||
"type": "string",
|
||
"description": "City name"
|
||
}
|
||
},
|
||
"required": ["city"]
|
||
}
|
||
}
|
||
]
|
||
)JSON";
|
||
|
||
const char* tool_response_tag = "tool_result";
|
||
|
||
// Mode select: "internal" (single run) or default "external" (two runs)
|
||
bool use_internal = (argc > 3 && std::string(argv[3]) == "internal");
|
||
|
||
if (aml_llm_enable_function_calling(
|
||
ctx,
|
||
/*system_prompt=*/nullptr,
|
||
tools_json,
|
||
tool_response_tag) != AML_LLM_Status_Success) {
|
||
std::cerr << "Enable function calling failed: aml_llm_enable_function_calling\n";
|
||
aml_llm_uninit(ctx);
|
||
return -1;
|
||
}
|
||
|
||
AML_LLMInput user_input;
|
||
std::memset(&user_input, 0, sizeof(user_input));
|
||
user_input.input_type = AML_LLM_INPUT_PROMPT;
|
||
user_input.prompt_input = "What's the weather in Hangzhou now?";
|
||
|
||
AML_LLMRunConfig run_cfg;
|
||
std::memset(&run_cfg, 0, sizeof(run_cfg));
|
||
run_cfg.run_mode = AML_LLM_RUN_GENERATE;
|
||
run_cfg.retain_history = 0;
|
||
|
||
if (use_internal) {
|
||
std::cout << "[Mode] Internal tool callback (single run)\n";
|
||
if (aml_llm_register_tool_callback(ctx, tool_callback, nullptr) != AML_LLM_Status_Success) {
|
||
std::cerr << "Register tool callback failed: aml_llm_register_tool_callback\n";
|
||
aml_llm_uninit(ctx);
|
||
return -1;
|
||
}
|
||
|
||
std::cout << "[User] " << user_input.prompt_input << "\n[Assistant] ";
|
||
DemoUserData ud; ud.printed = true;
|
||
AML_LLMRetStatus st = aml_llm_run(ctx, &user_input, &run_cfg, &ud);
|
||
if (st != AML_LLM_Status_Success) {
|
||
std::cerr << "\naml_llm_run failed\n";
|
||
}
|
||
|
||
} else {
|
||
std::cout << "[Mode] External tool execution (two runs)\n";
|
||
// First pass: ask the question and collect model output containing <tool_call>{...}
|
||
std::cout << "[User] " << user_input.prompt_input << "\n[Model] Planning tool calls...\n";
|
||
DemoUserData ud1; ud1.printed = true;
|
||
AML_LLMRetStatus st1 = aml_llm_run(ctx, &user_input, &run_cfg, &ud1);
|
||
if (st1 != AML_LLM_Status_Success) {
|
||
std::cerr << "First run aml_llm_run failed\n";
|
||
aml_llm_uninit(ctx);
|
||
return -1;
|
||
}
|
||
|
||
// Extract tool calls
|
||
auto calls = parse_tool_calls(ud1.buffer);
|
||
if (calls.empty()) {
|
||
std::cerr << "No tool calls parsed. First run output:\n" << ud1.buffer << "\n";
|
||
aml_llm_uninit(ctx);
|
||
return -1;
|
||
}
|
||
|
||
// Execute tools externally and build JSON array
|
||
std::string tool_results = build_tool_results_json(calls);
|
||
std::cout << "[External] Tool execution complete, results: " << tool_results << "\n";
|
||
|
||
// Second pass: feed tool results as prompt; SDK stitches full second-run prompt
|
||
|
||
|
||
AML_LLMInput tool_input;
|
||
std::memset(&tool_input, 0, sizeof(tool_input));
|
||
tool_input.input_type = AML_LLM_INPUT_PROMPT;
|
||
|
||
// ===================================================
|
||
// 2️⃣ External Mode: Second run after tool execution
|
||
// ===================================================
|
||
std::string system_prompt =
|
||
"You are a helpful assistant. "
|
||
"Use the provided tool results to answer the user's question naturally.";
|
||
|
||
// prefix: optional tool results block
|
||
std::string prompt_prefix;
|
||
prompt_prefix += "<|im_start|>user\n";
|
||
prompt_prefix += user_input.prompt_input; // e.g. "What's the weather in Hangzhou now?"
|
||
prompt_prefix += "\n<|im_end|>\n";
|
||
prompt_prefix += "<|im_start|>tool\n";
|
||
prompt_prefix += tool_results; // e.g. [{"city":"Hangzhou","temperature":"25C","condition":"Sunny"}]
|
||
prompt_prefix += "\n<|im_end|>\n<|im_start|>assistant\n";
|
||
|
||
// postfix: normally empty for direct answer
|
||
std::string prompt_postfix = "";
|
||
|
||
// // set a new chat template for answer mode
|
||
aml_llm_set_chat_template(
|
||
ctx,
|
||
system_prompt.c_str(),
|
||
prompt_prefix.c_str(),
|
||
prompt_postfix.c_str()
|
||
);
|
||
|
||
// prepare input
|
||
AML_LLMInput answer_input{};
|
||
answer_input.input_type = AML_LLM_INPUT_PROMPT;
|
||
answer_input.prompt_input = ""; // no user text, context is already in template
|
||
|
||
// run again
|
||
|
||
aml_llm_run(ctx, &answer_input, &run_cfg, &ud1);
|
||
// printf("Final Result: %s \n", ud2.buffer.c_str());
|
||
}
|
||
|
||
aml_llm_uninit(ctx);
|
||
return 0;
|
||
}
|