fix: add patch for search-release v0.3.0 and v0.4.0 to fiw uBatch value for model (#42)

* fix: add patch for search-release v0.3.0 and v0.4.0 to fiw uBatch value for model

Signed-off-by: ataldir <ataldir@cisco.com>

* chore: add a MAX_REQUEST_LEN parameter

Signed-off-by: ataldir <ataldir@cisco.com>

* chore: rewrite two conditions

Signed-off-by: ataldir <ataldir@cisco.com>

---------

Signed-off-by: ataldir <ataldir@cisco.com>

Author: thelasttoto

Makefile

@@ -70,6 +70,7 @@ search-release-$(SEARCH_VERSION)/README.md: download_models_for_semrouter
 	cd "search-release-$(SEARCH_VERSION)" && git lfs install --local
 	cd "search-release-$(SEARCH_VERSION)" && git submodule update --init --recursive
 	cd "search-release-$(SEARCH_VERSION)" && git lfs pull
+	cd "search-release-$(SEARCH_VERSION)" && git apply --ignore-space-change --ignore-whitespace --reject --quiet --whitespace=fix ../patches/0001-fix-ubatch-in-search-release-$(SEARCH_VERSION).patch  || true
 
 build_search_lib search-release-$(SEARCH_VERSION)/build/lib/$(SEARCH_LIB): search-release-$(SEARCH_VERSION)/README.md
 	-mkdir -p "search-release-$(SEARCH_VERSION)/build"

go.work.sum

@@ -253,6 +253,8 @@ github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
+github.com/kelindar/search v0.3.0 h1:NY6JkbC8Wy3vujwIR/p9NW3xaVDZxFrE+JQECiFzwW0=
+github.com/kelindar/search v0.3.0/go.mod h1:7goLXnzQ6b0vMJr9SKWmABblTrJgPG3rVwp3yz2Lo5Q=
 github.com/kelindar/search v0.4.0 h1:mj3U26qB6BQJr9/6Q1vHS/I40CQ6Mhb5iJfmfI2e/mY=
 github.com/kelindar/search v0.4.0/go.mod h1:7goLXnzQ6b0vMJr9SKWmABblTrJgPG3rVwp3yz2Lo5Q=
 github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=

patches/0001-fix-ubatch-in-search-release-v0.3.0.patch

@@ -0,0 +1,13 @@
+diff --git a/llama-go.cpp/llama-go.cpp b/llama-go.cpp/llama-go.cpp
+index 4670be1..53d2cec 100644
+--- a/llama-go.cpp/llama-go.cpp
++++ b/llama-go.cpp/llama-go.cpp
+@@ -100,6 +100,8 @@ extern "C" {
+         struct llama_context_params params = llama_context_default_params();
+         params.n_ctx = ctx_size;
+         params.embeddings = embeddings;
++        printf("[+] Set context param n_ubatch to 2048\n");
++        params.n_ubatch = 2048;
+         return llama_new_context_with_model(model, params);
+     }
+ 

patches/0001-fix-ubatch-in-search-release-v0.4.0.patch

@@ -0,0 +1,14 @@
+diff --git a/llama-go.cpp/llama-go.cpp b/llama-go.cpp/llama-go.cpp
+index 79c5ab3..c4e90d1 100644
+--- a/llama-go.cpp/llama-go.cpp
++++ b/llama-go.cpp/llama-go.cpp
+@@ -100,7 +100,8 @@ extern "C" {
+         struct llama_context_params params = llama_context_default_params();
+         params.n_ctx = ctx_size;
+         params.embeddings = embeddings; // Corrected field name
+-
++        printf("[+] Set context param n_ubatch to 2048\n");
++        params.n_ubatch = 2048;
+         return llama_init_from_model(model, params); // Updated function
+     }
+ 

plugins/agent_bridge.go

@@ -87,6 +87,12 @@ func SelectAndRewrite(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if apiConfig.MaxRequestLength > 0 && r.ContentLength > apiConfig.MaxRequestLength {
+		logger.Debugf("[+] Query is too large, ignoring ...")
+		http.Error(rw, "Query is too large", http.StatusRequestEntityTooLarge)
+		return
+	}
+
 	nlqBytes, err := io.ReadAll(r.Body)
 	if err != nil {
 		logger.Errorf("[+] Error while reading the body: %s", err)

plugins/agent_bridge_acp.go

@@ -33,6 +33,7 @@ type ACPPluginData struct {
 	ModelEmbedder     *search.Vectorizer
 	ModelIndex        *search.Index[string]
 	StoreVersion      int64
+	MaxRequestLength  int64 `json:"maxRequestLength"` // MaxRequestSize is the maximum size of the request in characters; default is -1 (no limit)
 }
 
 var acpPluginData = ACPPluginData{
@@ -64,6 +65,12 @@ func ProcessACPQuery(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if acpPluginData.MaxRequestLength > 0 && r.ContentLength > acpPluginData.MaxRequestLength {
+		logger.Debugf("[+] Query is too large, ignoring ...")
+		http.Error(rw, "Query is too large", http.StatusRequestEntityTooLarge)
+		return
+	}
+
 	nlqBytes, err := io.ReadAll(r.Body)
 	if err != nil {
 		logger.Errorf("[+] Error while reading the body: %s", err)
@@ -104,9 +111,13 @@ func ProcessACPQuery(rw http.ResponseWriter, r *http.Request) {
 func initACPPluginApiConfig() error {
 	// Clear existing map
 	acpPluginData.ACPPluginServices = make(map[string]ACPPluginApiConfig)
+
 	// save the current version of the store BEFORE retreiving the data
 	acpPluginData.StoreVersion = storeVersion
 	logger.Debugf("[+] Loading ACP plugin config version (%v) ...", acpPluginData.StoreVersion)
+
+	acpPluginData.MaxRequestLength = int64(getEnvAsInt("MAX_REQUEST_SIZE", DEFAULT_MAX_REQUEST_SIZE))
+
 	// Get All APIs keys and values from Redis
 	apiKeysValues := agentBridgeStore.GetKeysAndValuesWithFilter("*")
 	if apiKeysValues == nil {

plugins/agent_bridge_config.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strconv"
 	"sync"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai"
@@ -24,8 +25,9 @@ const (
 	DEFAULT_OPENAI_ENDPOINT        = "https://api.openai.com/v1"
 	DEFAULT_OPENAI_MODEL           = "gpt-4o-mini"
 
-	MAX_UTERANCE_LENGTH   = 1500
-	VECTORIZER_GPU_LAYERS = 1
+	MAX_UTERANCE_LENGTH      = 1500
+	VECTORIZER_GPU_LAYERS    = 1
+	DEFAULT_MAX_REQUEST_SIZE = -1 // in characters; -1 means no limit
 )
 
 type AzureConfig struct {
@@ -63,6 +65,8 @@ type PluginDataConfig struct {
 
 	APIID      string
 	ListenPath string
+
+	MaxRequestLength int64 `json:"maxRequestLength"` // MaxRequestSize is the maximum size of the request in characters; default is -1 (no limit)
 }
 
 func getApiId(r *http.Request) (string, error) {
@@ -102,6 +106,17 @@ func getEnvOrDefault(value string, envKey string, defaultValue string) string {
 	return value
 }
 
+func getEnvAsInt(envKey string, defaultValue int) int {
+	valueStr := os.Getenv(envKey)
+	if valueStr != "" {
+		if value, err := strconv.Atoi(valueStr); err == nil {
+			return value
+		}
+	}
+
+	return defaultValue
+}
+
 func parseConfigData(apiId string, configData map[string]any) (*PluginDataConfig, error) {
 	logger.Debugf("[+] Parsing config for api id: %s", apiId)
 
@@ -130,7 +145,8 @@ func parseConfigData(apiId string, configData map[string]any) (*PluginDataConfig
 		SelectModelsPath:     DEFAULT_MODEL_EMBEDDINGS_PATH,
 		RelevanceThreshold:   threshold,
 
-		APIID: apiId,
+		APIID:            apiId,
+		MaxRequestLength: int64(getEnvAsInt("MAX_REQUEST_SIZE", DEFAULT_MAX_REQUEST_SIZE)),
 	}
 
 	logger.Debugf("[+] Finished parsing config for api id: %s", apiId)

plugins/agent_bridge_endpoint_selection_test.go

@@ -166,3 +166,17 @@ func TestEndpointSelection(t *testing.T) {
 	}
 
 }
+
+func TestFixForUBatchBug(t *testing.T) {
+	pluginDataConfig, err := loadApiSpecsForTests("tyk-github-id", "../configs/api.github.com.gist.deref.oas.json")
+	assert.Nil(t, err)
+
+	modelPath := filepath.Join(pluginDataConfig.SelectModelsPath, DEFAULT_MODEL_EMBEDDINGS_MODEL)
+	modelEmbedder, err := search.NewVectorizer(modelPath, 1)
+	assert.Nil(t, err)
+
+	inputInFailure := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+	_, err = modelEmbedder.EmbedText(inputInFailure)
+	assert.Nil(t, err)
+
+}