Skip to content

Commit 1bf839b

Browse files
authored
Enhance user input handling for llama-run (#11138)
The main motivation for this change is it was not handing ctrl-c/ctrl-d correctly. Modify `read_user_input` to handle EOF, "/bye" command, and empty input cases. Introduce `get_user_input` function to manage user input loop and handle different return cases. Signed-off-by: Eric Curtin <ecurtin@redhat.com>
1 parent f7cd133 commit 1bf839b

File tree

1 file changed

+60
-3
lines changed

1 file changed

+60
-3
lines changed

examples/run/run.cpp

Lines changed: 60 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
# include <curl/curl.h>
1212
#endif
1313

14+
#include <signal.h>
15+
1416
#include <climits>
1517
#include <cstdarg>
1618
#include <cstdio>
@@ -25,6 +27,13 @@
2527
#include "json.hpp"
2628
#include "llama-cpp.h"
2729

30+
#if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__)) || defined(_WIN32)
31+
[[noreturn]] static void sigint_handler(int) {
32+
printf("\n");
33+
exit(0); // not ideal, but it's the only way to guarantee exit in all cases
34+
}
35+
#endif
36+
2837
GGML_ATTRIBUTE_FORMAT(1, 2)
2938
static std::string fmt(const char * fmt, ...) {
3039
va_list ap;
@@ -801,7 +810,20 @@ static int generate(LlamaData & llama_data, const std::string & prompt, std::str
801810

802811
static int read_user_input(std::string & user) {
803812
std::getline(std::cin, user);
804-
return user.empty(); // Should have data in happy path
813+
if (std::cin.eof()) {
814+
printf("\n");
815+
return 1;
816+
}
817+
818+
if (user == "/bye") {
819+
return 1;
820+
}
821+
822+
if (user.empty()) {
823+
return 2;
824+
}
825+
826+
return 0; // Should have data in happy path
805827
}
806828

807829
// Function to generate a response based on the prompt
@@ -868,15 +890,34 @@ static bool is_stdout_a_terminal() {
868890
#endif
869891
}
870892

871-
// Function to tokenize the prompt
893+
// Function to handle user input
894+
static int get_user_input(std::string & user_input, const std::string & user) {
895+
while (true) {
896+
const int ret = handle_user_input(user_input, user);
897+
if (ret == 1) {
898+
return 1;
899+
}
900+
901+
if (ret == 2) {
902+
continue;
903+
}
904+
905+
break;
906+
}
907+
908+
return 0;
909+
}
910+
911+
// Main chat loop function
872912
static int chat_loop(LlamaData & llama_data, const std::string & user) {
873913
int prev_len = 0;
874914
llama_data.fmtted.resize(llama_n_ctx(llama_data.context.get()));
875915
static const bool stdout_a_terminal = is_stdout_a_terminal();
876916
while (true) {
877917
// Get user input
878918
std::string user_input;
879-
while (handle_user_input(user_input, user)) {
919+
if (get_user_input(user_input, user) == 1) {
920+
return 0;
880921
}
881922

882923
add_message("user", user.empty() ? user_input : user, llama_data);
@@ -917,7 +958,23 @@ static std::string read_pipe_data() {
917958
return result.str();
918959
}
919960

961+
static void ctrl_c_handling() {
962+
#if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
963+
struct sigaction sigint_action;
964+
sigint_action.sa_handler = sigint_handler;
965+
sigemptyset(&sigint_action.sa_mask);
966+
sigint_action.sa_flags = 0;
967+
sigaction(SIGINT, &sigint_action, NULL);
968+
#elif defined(_WIN32)
969+
auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
970+
return (ctrl_type == CTRL_C_EVENT) ? (sigint_handler(SIGINT), true) : false;
971+
};
972+
SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
973+
#endif
974+
}
975+
920976
int main(int argc, const char ** argv) {
977+
ctrl_c_handling();
921978
Opt opt;
922979
const int ret = opt.init(argc, argv);
923980
if (ret == 2) {

0 commit comments

Comments
 (0)