refactor the substring search function

Author: matteoserva

tools/server/server.cpp

@@ -2216,19 +2216,12 @@ struct server_context {
                     search_pos = slot.generated_text.size() - search_len;
                 }
 
-                auto found_pos = slot.generated_text.npos;
-                bool found = false;
-                std::string found_string;
-                for (auto start_string: slot.params.start_strings) {
-                    found_pos = slot.generated_text.find(start_string,search_pos);
-                    if (found_pos != slot.generated_text.npos) {
-                        found = true;
-                        found_string = start_string;
-                        break;
-                    }
-                }
+                std::pair<size_t, int> search_result = find_first_substring(slot.generated_text,slot.params.start_strings, search_pos);
+                bool found = search_result.first != std::string::npos;
 
-                if (found && slot.generated_text.size() > (found_pos + found_string.size()) ) {
+                if (found) {
+                    auto found_pos = search_result.first;
+                    std::string found_string = slot.params.start_strings[search_result.second];
                     slot.generated_text.erase(
                         slot.generated_text.begin(),
                         slot.generated_text.begin() + found_pos + found_string.size());

tools/server/utils.hpp

@@ -489,6 +489,25 @@ static size_t find_partial_stop_string(const std::string &stop, const std::strin
     return std::string::npos;
 }
 
+/* returns a pair containing the position and index of first element found.
+   returns npos and -1 if not found */
+static std::pair<size_t, int> find_first_substring(const std::string &haystack, const std::vector<std::string> & needles, size_t search_pos = 0)
+{
+    size_t found_pos  = std::string::npos;
+    int    found_idx = -1;
+
+    for (unsigned int i = 0; i < needles.size(); ++i) {
+        const std::string & start_string = needles[i];
+        found_pos = haystack.find(start_string,search_pos);
+        if (found_pos != std::string::npos) {
+            found_idx = i;
+            break;
+        }
+    }
+
+    return std::pair<size_t, int>(found_pos,found_idx);
+}
+
 // TODO: reuse llama_detokenize
 template <class Iter>
 static std::string tokens_to_str(llama_context * ctx, Iter begin, Iter end) {