feat: integrate Google semantic search and branded desktop launcher

crates/searchi-semantic/native/worker.cpp

@@ -0,0 +1,221 @@
+// Local NDJSON helper around the pinned llama.cpp encoder. stdout is exclusively
+// protocol output; llama diagnostics go to stderr. No model download or HTTP.
+#include "llama.h"
+#include "nlohmann/json.hpp"
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <cstdio>
+#include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+using json = nlohmann::json;
+constexpr size_t max_request_bytes = 2 * 1024 * 1024;
+constexpr size_t max_passage_bytes = 1024 * 1024;
+constexpr size_t max_text_bytes = 65536;
+constexpr size_t max_windows = 256;
+constexpr int context_tokens = 2048;
+constexpr int dimensions = 768;
+
+static void require(bool condition, const char * message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+struct InferenceError : std::runtime_error { using std::runtime_error::runtime_error; };
+static void require_inference(bool condition, const char * message) {
+    if (!condition) throw InferenceError(message);
+}
+
+class Encoder {
+    llama_model * model = nullptr;
+    llama_context * context = nullptr;
+    llama_batch_ext * batch = nullptr;
+    const llama_vocab * vocab = nullptr;
+    std::string model_path;
+    void release() {
+        if (batch) llama_batch_ext_free(batch);
+        if (context) llama_free(context);
+        if (model) llama_model_free(model);
+        batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr;
+    }
+    bool open(bool gpu) {
+        auto mp = llama_model_default_params();
+        mp.n_gpu_layers = gpu ? 999 : 0;
+        mp.progress_callback = [](float,void *) {return true;};
+        model = llama_model_load_from_file(model_path.c_str(),mp);
+        if (!model) return false;
+        auto cp = llama_context_default_params();
+        cp.n_ctx = context_tokens; cp.n_batch = context_tokens; cp.n_ubatch = context_tokens;
+        cp.n_seq_max = 1; cp.kv_unified = true; cp.embeddings = true;
+        cp.pooling_type = LLAMA_POOLING_TYPE_MEAN;
+        cp.n_threads = threads; cp.n_threads_batch = threads;
+        context = llama_init_from_model(model,cp);
+        if (!context) {release(); return false;}
+        require(llama_model_n_embd_out(model) == dimensions,"EmbeddingGemma 2 output must have 768 dimensions");
+        vocab = llama_model_get_vocab(model);
+        batch = llama_batch_ext_init(context);
+        if (!batch) {release(); return false;}
+        device = gpu ? "gpu" : "cpu";
+        return true;
+    }
+public:
+    std::string device;
+    int threads;
+    Encoder(const std::string & path, int cpu_threads, bool want_gpu) : model_path(path),threads(cpu_threads) {
+        llama_log_set([](ggml_log_level level, const char * text, void *) {
+            if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
+        },nullptr);
+        llama_backend_init();
+        bool use_gpu = want_gpu && llama_supports_gpu_offload();
+        require(open(use_gpu) || (use_gpu && open(false)),"Cannot initialize local EmbeddingGemma 2 model/context");
+    }
+    ~Encoder() {
+        release();
+        llama_backend_free();
+    }
+    bool healthy() const {return context != nullptr && batch != nullptr;}
+    void retry_cpu() {
+        require(device == "gpu","Only a GPU inference failure can request a CPU retry");
+        release();
+        require(open(false),"CPU fallback initialization failed after GPU inference failure");
+    }
+    std::vector<llama_token> tokenize(const std::string & text, bool specials) const {
+        std::vector<llama_token> tokens(text.size() + 8);
+        int count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
+        if (count < 0) {
+            tokens.resize(size_t(-count));
+            count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
+        }
+        require(count >= 0, "Tokenizer failed");
+        tokens.resize(size_t(count));
+        return tokens;
+    }
+    std::string decode(const std::vector<llama_token> & tokens, size_t count) const {
+        std::vector<char> text(max_passage_bytes + 64);
+        int length = llama_detokenize(vocab, tokens.data(), int(count), text.data(), int(text.size()), false, false);
+        require(length >= 0, "Detokenized passage exceeds byte budget");
+        return std::string(text.data(),size_t(length));
+    }
+    static std::string document_prefix(const std::string & title) {
+        require(title.size() <= 512, "Document title exceeds 512 UTF-8 bytes");
+        return "title: " + (title.empty() ? "none" : title) + " | text: ";
+    }
+    std::vector<float> embed(const std::string & text, int cap) {
+        require(!text.empty() && text.size() <= max_text_bytes, "Embedding text must be nonempty and at most 64 KiB");
+        auto tokens = tokenize(text,true);
+        require(!tokens.empty() && tokens.size() <= size_t(cap), "Prompt exceeds model token budget; split document windows first");
+        llama_batch_ext_clear(batch);
+        for (size_t i = 0; i < tokens.size(); ++i) {
+            int index = llama_batch_ext_add_token(batch,0,tokens[i]);
+            require(index >= 0, "Embedding batch is full");
+            llama_pos position[4] = {llama_pos(i),0,0,0};
+            require(llama_batch_ext_set_pos(batch,index,position), "Cannot set encoder position");
+            require(llama_batch_ext_set_output_embd(batch,index,true), "Cannot select embedding output");
+        }
+        require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch) == 0, "llama_process encoder failed");
+        const float * pooled = llama_get_embeddings_seq(context,0);
+        require_inference(pooled != nullptr, "Encoder returned no pooled vector");
+        double norm = 0;
+        for (int i = 0; i < dimensions; ++i) { require_inference(std::isfinite(pooled[i]), "Nonfinite encoder output"); norm += double(pooled[i])*pooled[i]; }
+        require_inference(norm > 0 && std::isfinite(norm), "Zero or nonfinite embedding norm");
+        std::vector<float> vector(dimensions);
+        for (int i = 0; i < dimensions; ++i) vector[i] = float(pooled[i]/std::sqrt(norm));
+        return vector;
+    }
+    json windows(const std::string & text, const std::string & title) const {
+        require(!text.empty() && text.size() <= max_passage_bytes, "Passage exceeds 1 MiB or is empty");
+        auto tokens = tokenize(text,false);
+        require(!tokens.empty(), "Passage has no tokens");
+        // Preserve actual source offsets, rather than inventing locations from
+        // a lossy decoded representation. Unexpected normalization is an error.
+        require(decode(tokens,tokens.size()) == text, "Tokenizer roundtrip changed source text; cannot verify source offsets");
+        size_t prompt = tokenize(document_prefix(title),true).size();
+        require(prompt + 132 < size_t(context_tokens), "Document prefix exceeds token budget");
+        size_t budget = size_t(context_tokens) - prompt - 4;
+        size_t stride = budget - 128;
+        json result = json::array();
+        for (size_t start = 0; start < tokens.size(); start += stride) {
+            require(result.size() < max_windows, "Passage exceeds 256 windows");
+            size_t end = std::min(start + budget,tokens.size());
+            size_t begin_byte = start ? decode(tokens,start).size() : 0;
+            size_t end_byte = decode(tokens,end).size();
+            // Token pieces can split a UTF-8 character. Enlarge to valid source
+            // boundaries; the four-token reserve covers boundary re-tokenization.
+            while (begin_byte > 0 && begin_byte < text.size() && (uint8_t(text[begin_byte]) & 0xc0) == 0x80) --begin_byte;
+            while (end_byte < text.size() && (uint8_t(text[end_byte]) & 0xc0) == 0x80) ++end_byte;
+            require(begin_byte < end_byte && end_byte <= text.size(), "Invalid source window byte range");
+            auto fragment = text.substr(begin_byte,end_byte-begin_byte);
+            require(fragment.size() + document_prefix(title).size() <= max_text_bytes, "Window exceeds text byte budget");
+            require(tokenize(document_prefix(title)+fragment,true).size() <= size_t(context_tokens), "Window re-tokenization exceeds context budget");
+            result.push_back({{"text",fragment},{"start_byte",begin_byte},{"end_byte",end_byte}});
+            if (end == tokens.size()) break;
+        }
+        return result;
+    }
+    json handle(const json & request) {
+        auto op = request.at("op").get<std::string>();
+        if (op == "windows") return {{"windows",windows(request.at("text").get<std::string>(),request.value("title",std::string()))}};
+        if (op == "query") return {{"vectors",json::array({embed("task: search result | query: "+request.at("text").get<std::string>(),512)})}};
+        if (op == "documents") {
+            const auto & texts = request.at("texts");
+            require(texts.is_array() && !texts.empty() && texts.size() <= 8,"Batch size must be between 1 and 8");
+            auto prefix = document_prefix(request.value("title",std::string()));
+            json vectors = json::array();
+            for (const auto & text : texts) vectors.push_back(embed(prefix+text.get<std::string>(),context_tokens));
+            return {{"vectors",vectors}};
+        }
+        throw std::runtime_error("Unknown embedding operation");
+    }
+};
+
+int main(int argc, char ** argv) {
+    try {
+        std::string path, device = "auto"; int threads = 1;
+        for (int i = 1; i < argc; ++i) {
+            require(i+1 < argc,"Worker options require values");
+            std::string option = argv[i++], value = argv[i];
+            if (option == "--model") path = value;
+            else if (option == "--threads") threads = std::stoi(value);
+            else if (option == "--device") device = value;
+            else throw std::runtime_error("Unknown worker option");
+        }
+        require(!path.empty() && threads >= 1 && threads <= 8,"Model path required; threads must be 1–8");
+        require(device == "cpu" || device == "auto","Device must be cpu or auto");
+        bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
+        Encoder encoder(path,threads,device == "auto" && !force_cpu);
+        std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
+            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG}}).dump() << '\n' << std::flush;
+        for (;;) {
+            std::string line; char c;
+            while (std::cin.get(c) && c != '\n') {
+                require(line.size() < max_request_bytes,"Request exceeds 2 MiB"); line.push_back(c);
+            }
+            if (!std::cin && line.empty()) break;
+            json response;
+            try {
+                auto request = json::parse(line);
+                try {response = encoder.handle(request);}
+                catch (const InferenceError &) {
+                    if (encoder.device != "gpu") throw;
+                    encoder.retry_cpu();
+                    response = encoder.handle(request); // Exactly one CPU retry.
+                    response["fallback"] = "gpu_inference_failure";
+                }
+                response["ok"] = true;
+            }
+            catch (const std::exception & error) { response = {{"ok",false},{"error",error.what()},{"fatal",!encoder.healthy()}}; }
+            response["device"] = encoder.device;
+            auto encoded = response.dump();
+            require(encoded.size() <= 4*1024*1024,"Response exceeds 4 MiB");
+            std::cout << encoded << '\n' << std::flush;
+            if (!encoder.healthy()) return 1;
+        }
+        return 0;
+    } catch (const std::exception & error) {
+        std::cerr << "Searchi Gemma worker: " << error.what() << '\n';
+        return 1;
+    }
+}