mudler
diff --git a/‎.github/workflows/test-extra.yml‎
Lines changed: 1 addition & 0 deletions b/‎.github/workflows/test-extra.yml‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎.github/workflows/test.yml‎
Lines changed: 6 additions & 0 deletions b/‎.github/workflows/test.yml‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎Makefile‎
Lines changed: 12 additions & 0 deletions b/‎Makefile‎
Lines changed: 12 additions & 0 deletions
diff --git a/‎backend/backend.proto‎
Lines changed: 20 additions & 0 deletions b/‎backend/backend.proto‎
Lines changed: 20 additions & 0 deletions
diff --git a/‎backend/cpp/ik-llama-cpp/grpc-server.cpp‎
Lines changed: 1 addition & 0 deletions b/‎backend/cpp/ik-llama-cpp/grpc-server.cpp‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎backend/cpp/llama-cpp/grpc-server.cpp‎
Lines changed: 28 additions & 30 deletions b/‎backend/cpp/llama-cpp/grpc-server.cpp‎
Lines changed: 28 additions & 30 deletions
diff --git a/‎backend/go/local-store/store.go‎
Lines changed: 7 additions & 2 deletions b/‎backend/go/local-store/store.go‎
Lines changed: 7 additions & 2 deletions
diff --git a/‎backend/go/local-store/store_test.go‎
Lines changed: 40 additions & 0 deletions b/‎backend/go/local-store/store_test.go‎
Lines changed: 40 additions & 0 deletions
diff --git a/‎backend/python/insightface/backend.py‎
Lines changed: 4 additions & 1 deletion b/‎backend/python/insightface/backend.py‎
Lines changed: 4 additions & 1 deletion
diff --git a/‎backend/python/tinygrad/backend.py‎
Lines changed: 4 additions & 1 deletion b/‎backend/python/tinygrad/backend.py‎
Lines changed: 4 additions & 1 deletion
@@ -526,6 +526,7 @@ jobs:
       - name: Build llama-cpp backend image and run gRPC e2e tests
         run: |
           make test-extra-backend-llama-cpp
+          make test-extra-backend-llama-cpp-embeddings
   tests-llama-cpp-grpc-transcription:
     needs: detect-changes
     if: needs.detect-changes.outputs.llama-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
 
@@ -65,6 +65,12 @@ jobs:
       - name: Test (with coverage gate)
         run: |
           PATH="$PATH:/root/go/bin" make --jobs 5 --output-sync=target test-coverage-check
+      # tests/integration is outside the coverage roots because its store specs
+      # need a live backend. test-stores builds and installs local-store before
+      # running the complete suite, so new local-store specs are collected
+      # automatically without adding another workflow entry.
+      - name: Test local-store integration
+        run: PATH="$PATH:$HOME/go/bin" make test-stores
       - name: Upload coverage report
         if: ${{ always() }}
         uses: actions/upload-artifact@v4
 
@@ -676,6 +676,7 @@ test-extra: prepare-test-extra
 ##   BACKEND_TEST_PROMPT      Override the prompt used in predict/stream specs.
 ##   BACKEND_TEST_OPTIONS     Comma-separated Options[] entries forwarded to LoadModel,
 ##                            e.g. "tool_parser:hermes,reasoning_parser:qwen3".
+##   BACKEND_TEST_EMBEDDING_LAYOUT Expected EmbeddingResult layout: "final" or "per_token".
 ##
 ## Direct usage (image already built, no docker-build-* dependency):
 ##
@@ -705,6 +706,7 @@ test-extra-backend: protogen-go
 	BACKEND_TEST_CAPS="$$BACKEND_TEST_CAPS" \
 	BACKEND_TEST_PROMPT="$$BACKEND_TEST_PROMPT" \
 	BACKEND_TEST_OPTIONS="$$BACKEND_TEST_OPTIONS" \
+	BACKEND_TEST_EMBEDDING_LAYOUT="$$BACKEND_TEST_EMBEDDING_LAYOUT" \
 	BACKEND_TEST_TOOL_PROMPT="$$BACKEND_TEST_TOOL_PROMPT" \
 	BACKEND_TEST_TOOL_NAME="$$BACKEND_TEST_TOOL_NAME" \
 	BACKEND_TEST_CACHE_TYPE_K="$$BACKEND_TEST_CACHE_TYPE_K" \
@@ -724,6 +726,15 @@ test-extra-backend-llama-cpp: docker-build-llama-cpp
 	BACKEND_TEST_CAPS=health,load,predict,stream,logprobs,logit_bias \
 	$(MAKE) test-extra-backend
 
+## Raw llama.cpp embeddings are required by Go-side pooling. This exercises the
+## real C++ backend and verifies that it marks the flattened matrix per-token.
+test-extra-backend-llama-cpp-embeddings: docker-build-llama-cpp
+	BACKEND_IMAGE=local-ai-backend:llama-cpp \
+	BACKEND_TEST_CAPS=health,load,embeddings \
+	BACKEND_TEST_OPTIONS=pooling:none \
+	BACKEND_TEST_EMBEDDING_LAYOUT=per_token \
+	$(MAKE) test-extra-backend
+
 test-extra-backend-ik-llama-cpp: docker-build-ik-llama-cpp
 	BACKEND_IMAGE=local-ai-backend:ik-llama-cpp $(MAKE) test-extra-backend
 
@@ -813,6 +824,7 @@ test-extra-backend-tinygrad-embeddings: docker-build-tinygrad
 	BACKEND_IMAGE=local-ai-backend:tinygrad \
 	BACKEND_TEST_MODEL_NAME=Qwen/Qwen3-0.6B \
 	BACKEND_TEST_CAPS=health,load,embeddings \
+	BACKEND_TEST_EMBEDDING_LAYOUT=final \
 	$(MAKE) test-extra-backend
 
 ## tinygrad — Stable Diffusion 1.5. The original CompVis/runwayml repos have
 
@@ -536,8 +536,28 @@ message Result {
   bool success = 2;
 }
 
+// EmbeddingLayout describes whether embeddings contains one final vector or
+// a matrix of per-token vectors. Go-side pooling must never infer this from
+// tokens/dim alone: a one-token raw matrix and a final vector have the same
+// shape.
+enum EmbeddingLayout {
+  EMBEDDING_LAYOUT_UNSPECIFIED = 0;
+  EMBEDDING_LAYOUT_FINAL = 1;
+  EMBEDDING_LAYOUT_PER_TOKEN = 2;
+}
+
 message EmbeddingResult {
   repeated float embeddings = 1;
+  // Shape of the payload above: dim is the embedding width, tokens is the
+  // number of vectors packed into `embeddings` (1 when the backend pooled
+  // server-side, N with pooling:none; total across prompts if a request
+  // carried several). tokens=0/dim=0 means the backend predates shape
+  // reporting. prompt_tokens is the number of prompt tokens evaluated, for
+  // usage accounting.
+  int32 tokens = 2;
+  int32 dim = 3;
+  int32 prompt_tokens = 4;
+  EmbeddingLayout layout = 5;
 }
 
 message TranscriptRequest {
 
@@ -2565,6 +2565,7 @@ class BackendServiceImpl final : public backend::Backend::Service {
     grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) {
         auto identity = checkModelIdentity(request);
         if (!identity.ok()) return identity;
+        embeddingResult->set_layout(backend::EMBEDDING_LAYOUT_FINAL);
         json data = parse_options(false, request, llama);
         const int task_id = llama.queue_tasks.get_new_id();
         llama.queue_results.add_waiting_task_id(task_id);
 
@@ -2911,42 +2911,40 @@ class BackendServiceImpl final : public backend::Backend::Service {
             return grpc::Status(grpc::StatusCode::INTERNAL, all_results.error->to_json().value("message", "Error in receiving results"));
         }
 
-        // Collect responses
-        json responses = json::array();
+        // Extract the embeddings typed, straight from the task results (no
+        // JSON round-trip), and report the payload shape alongside the same
+        // flat float array as before: dim is the embedding width, tokens the
+        // number of vectors packed into `embeddings` (1 per prompt when the
+        // server pooled, one per token with pooling:none; summed across
+        // prompts if the request carried several), prompt_tokens the prompt
+        // tokens evaluated, for usage accounting. Consumers seeing 0/0 know
+        // the backend predates shape reporting.
+        int32_t n_vectors = 0;
+        int32_t dim = 0;
+        int32_t prompt_tokens = 0;
         for (auto & res : all_results.results) {
-            GGML_ASSERT(dynamic_cast<server_task_result_embd*>(res.get()) != nullptr);
-            responses.push_back(res->to_json());
-        }
-
-        std::cout << "[DEBUG] Responses size: " << responses.size() << std::endl;
-
-        // Process the responses and extract embeddings
-        for (const auto & response_elem : responses) {
-            // Check if the response has an "embedding" field
-            if (response_elem.contains("embedding")) {
-                json embedding_data = json_value(response_elem, "embedding", json::array());
-
-                if (embedding_data.is_array() && !embedding_data.empty()) {
-                    for (const auto & embedding_vector : embedding_data) {
-                        if (embedding_vector.is_array()) {
-                            for (const auto & embedding_value : embedding_vector) {
-                                embeddingResult->add_embeddings(embedding_value.get<float>());
-                            }
-                        }
-                    }
+            auto * embd_res = dynamic_cast<server_task_result_embd*>(res.get());
+            GGML_ASSERT(embd_res != nullptr);
+            prompt_tokens += embd_res->n_tokens;
+            for (const auto & vec : embd_res->embedding) {
+                for (const float value : vec) {
+                    embeddingResult->add_embeddings(value);
                 }
-            } else {
-                // Check if the response itself contains the embedding data directly
-                if (response_elem.is_array()) {
-                    for (const auto & embedding_value : response_elem) {
-                        embeddingResult->add_embeddings(embedding_value.get<float>());
-                    }
+                if (!vec.empty()) {
+                    n_vectors++;
+                    dim = (int32_t) vec.size();
                 }
             }
         }
+        embeddingResult->set_tokens(n_vectors);
+        embeddingResult->set_dim(dim);
+        embeddingResult->set_prompt_tokens(prompt_tokens);
+        embeddingResult->set_layout(
+            llama_pooling_type(ctx_server.get_llama_context()) == LLAMA_POOLING_TYPE_NONE
+                ? backend::EMBEDDING_LAYOUT_PER_TOKEN
+                : backend::EMBEDDING_LAYOUT_FINAL);
 
-
-
+        std::cout << "[DEBUG] Embedding vectors: " << n_vectors << " x " << dim << std::endl;
 
         return grpc::Status::OK;
     }
 
@@ -38,8 +38,9 @@ type Store struct {
 	// keysAreNormalized stays true until any non-unit-magnitude key
 	// is added; once false, the magnitude-aware fallback path is
 	// used by Find. Re-evaluated only at Set time, never again on
-	// its own — a deletion of the offending key does NOT flip it
-	// back to true (the bookkeeping cost would dominate the gain).
+	// its own — a partial deletion of the offending key does NOT flip
+	// it back to true (the bookkeeping cost would dominate the gain).
+	// An empty store returns to its initial state.
 	keysAreNormalized bool
 
 	// keyLen is the dimension of every stored key. -1 means "no
@@ -142,6 +143,10 @@ func (s *Store) StoresDelete(opts *pb.StoresDeleteOptions) error {
 	mergedV = append(mergedV, tailV...)
 	s.keys = mergedK
 	s.values = mergedV
+	if len(s.keys) == 0 {
+		s.keyLen = -1
+		s.keysAreNormalized = true
+	}
 	assert(slices.IsSortedFunc(s.keys, slices.Compare[[]float32]), "Delete: s.keys not sorted post-merge")
 	assert(len(s.keys) == len(s.values), "Delete: keys/values length skew")
 	return nil
 
@@ -105,6 +105,46 @@ var _ = Describe("StoresDelete", func() {
 		})).To(Succeed(), "delete of missing key should succeed")
 		Expect(s.keys).To(HaveLen(1))
 	})
+
+	It("reopens the dimension after deleting every key", func() {
+		s := NewStore()
+		oldKey := []float32{2, 0, 0}
+		mustSet(s, [][]float32{oldKey}, [][]byte{[]byte("3d")})
+		Expect(s.keysAreNormalized).To(BeFalse())
+
+		Expect(s.StoresDelete(&pb.StoresDeleteOptions{
+			Keys: wrapKeys([][]float32{oldKey}),
+		})).To(Succeed())
+		Expect(s.keys).To(BeEmpty())
+		Expect(s.keyLen).To(Equal(-1))
+		Expect(s.keysAreNormalized).To(BeTrue())
+
+		newKey := normalizeVec([]float32{1, 1})
+		mustSet(s, [][]float32{newKey}, [][]byte{[]byte("2d")})
+		res, err := s.StoresFind(&pb.StoresFindOptions{
+			Key:  &pb.StoresKey{Floats: newKey},
+			TopK: 1,
+		})
+		Expect(err).NotTo(HaveOccurred())
+		Expect(res.Values).To(HaveLen(1))
+		Expect(string(res.Values[0].Bytes)).To(Equal("2d"))
+	})
+
+	It("retains the dimension after a partial delete", func() {
+		s := NewStore()
+		mustSet(s,
+			[][]float32{{1, 0, 0}, {0, 1, 0}},
+			[][]byte{[]byte("x"), []byte("y")},
+		)
+		Expect(s.StoresDelete(&pb.StoresDeleteOptions{
+			Keys: wrapKeys([][]float32{{1, 0, 0}}),
+		})).To(Succeed())
+		Expect(s.keyLen).To(Equal(3))
+		Expect(s.StoresSet(&pb.StoresSetOptions{
+			Keys:   wrapKeys([][]float32{{1, 0}}),
+			Values: wrapValues([][]byte{[]byte("2d")}),
+		})).NotTo(Succeed())
+	})
 })
 
 var _ = Describe("StoresFind", func() {
 
@@ -127,7 +127,10 @@ def Embedding(self, request, context):
             context.set_code(grpc.StatusCode.NOT_FOUND)
             context.set_details("no face detected")
             return backend_pb2.EmbeddingResult()
-        return backend_pb2.EmbeddingResult(embeddings=[float(x) for x in vec])
+        return backend_pb2.EmbeddingResult(
+            embeddings=[float(x) for x in vec],
+            layout=backend_pb2.EMBEDDING_LAYOUT_FINAL,
+        )
 
     def Detect(self, request, context):
         if self.engine is None:
 
@@ -638,7 +638,10 @@ async def Embedding(self, request, context):
             normalized = (pooled / (norm + 1e-12))
             vec = normalized.cast(dtypes.float32).tolist()
 
-            return backend_pb2.EmbeddingResult(embeddings=[float(x) for x in vec])
+            return backend_pb2.EmbeddingResult(
+                embeddings=[float(x) for x in vec],
+                layout=backend_pb2.EMBEDDING_LAYOUT_FINAL,
+            )
         except Exception as exc:
             import traceback
             traceback.print_exc()