node-similarity-search-native/similarity_search_addon.cc

#include <napi.h>
#include <string>
#include "similarity_search.h"

class SearchIndexWrapper : public Napi::ObjectWrap<SearchIndexWrapper> {
public:
  static Napi::Object Init(Napi::Env env, Napi::Object exports);
  SearchIndexWrapper(const Napi::CallbackInfo& info);
  ~SearchIndexWrapper();

private:
  static Napi::FunctionReference constructor;
  
  Napi::Value AddString(const Napi::CallbackInfo& info);
  Napi::Value Search(const Napi::CallbackInfo& info);
  Napi::Value GetSize(const Napi::CallbackInfo& info);
  
  SearchIndex* index_;
};

Napi::FunctionReference SearchIndexWrapper::constructor;

Napi::Object SearchIndexWrapper::Init(Napi::Env env, Napi::Object exports) {
  Napi::HandleScope scope(env);
  
  Napi::Function func = DefineClass(env, "SearchIndex", {
    InstanceMethod("addString", &SearchIndexWrapper::AddString),
    InstanceMethod("search", &SearchIndexWrapper::Search),
    InstanceMethod("size", &SearchIndexWrapper::GetSize)
  });
  
  constructor = Napi::Persistent(func);
  constructor.SuppressDestruct();
  
  exports.Set("SearchIndex", func);
  return exports;
}

SearchIndexWrapper::SearchIndexWrapper(const Napi::CallbackInfo& info) 
  : Napi::ObjectWrap<SearchIndexWrapper>(info) {
  Napi::Env env = info.Env();
  Napi::HandleScope scope(env);
  
  int capacity = 10000; // Increased default capacity from 500 to 10000
  if (info.Length() > 0 && info[0].IsNumber()) {
    capacity = info[0].As<Napi::Number>().Int32Value();
  }
  
  this->index_ = create_search_index(capacity);
  if (!this->index_) {
    Napi::Error::New(env, "Failed to create search index").ThrowAsJavaScriptException();
  }
}

SearchIndexWrapper::~SearchIndexWrapper() {
  free_search_index(this->index_);
}

Napi::Value SearchIndexWrapper::AddString(const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();
  Napi::HandleScope scope(env);
  
  if (info.Length() < 1 || !info[0].IsString()) {
    Napi::TypeError::New(env, "String expected").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  std::string str = info[0].As<Napi::String>().Utf8Value();
  
  // Check if string is empty
  if (str.empty()) {
    Napi::Error::New(env, "Empty string not allowed").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  // Check if string is too long
  if (str.length() >= MAX_STRING_LEN) {
    Napi::Error::New(env, "String too long").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  // Check if we've reached capacity
  if (this->index_->num_strings >= this->index_->capacity) {
    Napi::Error::New(env, "Search index capacity exceeded").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  int result = add_string_to_index(this->index_, str.c_str());
  if (result != 0) {
    Napi::Error::New(env, "Failed to add string to index").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  return Napi::Number::New(env, result);
}

Napi::Value SearchIndexWrapper::Search(const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();
  Napi::HandleScope scope(env);
  
  if (info.Length() < 1 || !info[0].IsString()) {
    Napi::TypeError::New(env, "Query string expected").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  std::string query = info[0].As<Napi::String>().Utf8Value();
  
  // Check if query is empty
  if (query.empty()) {
    Napi::Error::New(env, "Empty query not allowed").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  // Check if query string is too long
  if (query.length() >= MAX_STRING_LEN) {
    Napi::Error::New(env, "Query string too long").ThrowAsJavaScriptException();
    return env.Null();
  }
  
  float cutoff = 0.2f; // Default cutoff
  if (info.Length() > 1 && info[1].IsNumber()) {
    cutoff = info[1].As<Napi::Number>().FloatValue();
    if (cutoff < 0.0f || cutoff > 1.0f) {
      Napi::Error::New(env, "Cutoff must be between 0 and 1").ThrowAsJavaScriptException();
      return env.Null();
    }
  }
  
  int num_results = 0;
  SearchResult* results = search_index(this->index_, query.c_str(), cutoff, &num_results);
  
  // If no results found, return empty array instead of throwing error
  if (!results || num_results == 0) {
    return Napi::Array::New(env, 0);
  }
  
  Napi::Array result_array = Napi::Array::New(env, num_results);
  
  for (int i = 0; i < num_results; i++) {
    Napi::Object obj = Napi::Object::New(env);
    obj.Set("string", Napi::String::New(env, results[i].string));
    obj.Set("similarity", Napi::Number::New(env, results[i].similarity));
    result_array[i] = obj;
  }
  
  free_search_results(results, num_results);
  
  return result_array;
}

Napi::Value SearchIndexWrapper::GetSize(const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();
  Napi::HandleScope scope(env);
  
  return Napi::Number::New(env, this->index_->num_strings);
}

Napi::Object Init(Napi::Env env, Napi::Object exports) {
  return SearchIndexWrapper::Init(env, exports);
}

NODE_API_MODULE(similarity_search_addon, Init)