diff --git a/.github/workflows/test_npu_embedder.yml b/.github/workflows/test_npu_embedder.yml new file mode 100644 index 000000000..9e5a4707f --- /dev/null +++ b/.github/workflows/test_npu_embedder.yml @@ -0,0 +1,88 @@ +# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +# Validates the NPU-native FLM embedder (embed-gemma-300m-FLM, #1744) on real +# Ryzen AI hardware. The test (tests/test_npu_flm_embedder_hw.py) SKIPS unless +# the live Lemonade Server advertises the FLM model, so on a runner without the +# FLM/NPU backend this job is a clean no-op rather than a false failure. +name: Test NPU FLM Embedder + +on: + workflow_dispatch: + push: + branches: [ main ] + paths: + - 'src/gaia/agents/registry.py' + - 'src/gaia/agents/base/memory.py' + - 'src/gaia/agents/base/memory_store.py' + - 'src/gaia/llm/lemonade_client.py' + - 'tests/test_npu_flm_embedder_hw.py' + - '.github/workflows/test_npu_embedder.yml' + pull_request: + branches: [ main ] + types: [opened, synchronize, reopened, ready_for_review] + paths: + - 'src/gaia/agents/registry.py' + - 'src/gaia/agents/base/memory.py' + - 'src/gaia/agents/base/memory_store.py' + - 'src/gaia/llm/lemonade_client.py' + - 'tests/test_npu_flm_embedder_hw.py' + - '.github/workflows/test_npu_embedder.yml' + workflow_call: + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + test-npu-embedder: + name: Test NPU FLM Embedder + runs-on: [self-hosted, strix-halo] + if: github.event_name != 'pull_request' || github.event.pull_request.draft == false || contains(github.event.pull_request.labels.*.name, 'ready_for_ci') + + steps: + - uses: actions/checkout@v6 + + - name: Setup Python environment + uses: ./.github/actions/setup-venv + with: + python-version: '3.12' + install-package: '.[dev,rag]' + extra-index-url: 'https://download.pytorch.org/whl/cpu' + + - name: Install Lemonade Server + uses: ./.github/actions/install-lemonade + + - name: Start Lemonade + validate FLM embedder on NPU + shell: powershell + env: + NO_PROXY: "localhost,127.0.0.1" + run: | + [Console]::OutputEncoding = [System.Text.Encoding]::UTF8 + $env:PYTHONIOENCODING = "utf-8" + $env:PYTHONUTF8 = "1" + chcp 65001 | Out-Null + + # Free the port, then start Lemonade in a single session so the FLM + # backend child survives for the test run. + .\installer\scripts\start-lemonade.ps1 ` + -ModelName "gemma4-it-e2b-FLM" ` + -Port 13305 ` + -CtxSize 32768 ` + -InitWaitTime 15 + + try { + # The test loads embed-gemma-300m-FLM itself and skips cleanly if + # the runner's Lemonade has no FLM/NPU backend. -rs prints skip + # reasons so a silent skip is visible in the log. + python -m pytest tests/test_npu_flm_embedder_hw.py -v -rs --tb=short + if ($LASTEXITCODE -ne 0) { throw "NPU FLM embedder tests failed" } + } + finally { + if ($env:LEMONADE_PROCESS_ID) { + Stop-Process -Id $env:LEMONADE_PROCESS_ID -Force -ErrorAction SilentlyContinue + } + } diff --git a/hub/agents/python/chat/gaia_agent_chat/agent.py b/hub/agents/python/chat/gaia_agent_chat/agent.py index 1ff3ed777..e1905374b 100644 --- a/hub/agents/python/chat/gaia_agent_chat/agent.py +++ b/hub/agents/python/chat/gaia_agent_chat/agent.py @@ -21,7 +21,7 @@ from gaia.agents.base.agent import Agent, default_max_steps from gaia.agents.base.console import AgentConsole -from gaia.agents.base.memory import EMBEDDING_MODEL, MemoryMixin +from gaia.agents.base.memory import MemoryMixin # dynamic_tools_env_override is re-exported so callers importing it from # gaia_agent_chat.agent keep working; its canonical home is the core tool_loader @@ -29,6 +29,7 @@ from gaia.agents.base.tool_loader import dynamic_tools_env_override # noqa: F401 from gaia.agents.base.tool_loader import ToolLoader from gaia.agents.base.tools import _TOOL_REGISTRY +from gaia.agents.registry import get_embedding_model_for_device from gaia.agents.tools import FileSystemToolsMixin # Enhanced file system navigation from gaia.agents.tools import ScratchpadToolsMixin # Structured data analysis from gaia.agents.tools import ( # Web browsing and search; Shared tools @@ -259,10 +260,17 @@ def __init__(self, config: Optional[ChatAgentConfig] = None): else os.getenv("LEMONADE_BASE_URL", "http://localhost:13305/api/v1") ) + # Embedder is device-scoped: the NPU profile uses the FLM-native + # embedder so chat and embeddings stay co-resident on the NPU backend + # (a GGUF embedder runs on Vulkan and evicts the FLM chat model every + # turn — #1744). GPU/CPU keep the GGUF nomic embedder. + effective_embedding_model = get_embedding_model_for_device(config.device) + # Initialize RAG SDK (optional - will be None if dependencies not installed) try: rag_config = RAGConfig( model=effective_model_id, + embedding_model=effective_embedding_model, chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap, # Configurable overlap for context preservation max_chunks=config.max_chunks, @@ -352,7 +360,7 @@ def __init__(self, config: Optional[ChatAgentConfig] = None): self.tool_loader = self._maybe_build_tool_loader() # Initialize memory subsystem (before super().__init__ which calls _register_tools) - self.init_memory() + self.init_memory(embedding_model=effective_embedding_model) # Store base URL for use in _register_tools() (VLM, etc.) self._base_url = effective_base_url @@ -503,7 +511,7 @@ def _embed_texts_batch(self, texts) -> "Any": """ import numpy as np - results = self._get_embedder().embed(list(texts), model=EMBEDDING_MODEL) + results = self._get_embedder().embed(list(texts), model=self._embedding_model) vecs = np.asarray(results, dtype=np.float32) norms = np.linalg.norm(vecs, axis=1, keepdims=True) norms[norms == 0] = 1.0 diff --git a/src/gaia/agents/base/memory.py b/src/gaia/agents/base/memory.py index 4d72f043a..606d52266 100644 --- a/src/gaia/agents/base/memory.py +++ b/src/gaia/agents/base/memory.py @@ -146,10 +146,15 @@ def _changed_software_versions(existing: List[Dict]) -> List[str]: # Constants # ============================================================================ -#: Embedding model served by Lemonade — 768-dim, MOE architecture. +#: Default embedder served by Lemonade — 768-dim GGUF (GPU/CPU profiles). +#: The active embedder is per-instance (``self._embedding_model``) and may be +#: the NPU-native FLM embedder instead; see ``init_memory`` (#1744). These +#: module constants remain the fallback default. EMBEDDING_MODEL = "nomic-embed-text-v2-moe-GGUF" -#: Embedding dimensionality for nomic-embed-text-v2-moe. +#: Default embedding dimensionality (nomic-embed-text-v2-moe). The active dim is +#: derived from the live embedder at startup (``self._embedding_dim``); this is +#: only the pre-probe fallback. EMBEDDING_DIM = 768 #: Cross-encoder model for reranking (~22 MB, runs on CPU). @@ -308,7 +313,10 @@ class MemoryMixin(ProceduralMemoryMixin): """ def init_memory( - self, db_path: Optional[Path] = None, context: str = "global" + self, + db_path: Optional[Path] = None, + context: str = "global", + embedding_model: Optional[str] = None, ) -> None: """Initialize the memory subsystem (v2 startup sequence). @@ -321,6 +329,11 @@ def init_memory( Args: db_path: Optional path for the DB file. Default: ~/.gaia/memory.db context: Active context scope (e.g., 'work', 'personal', 'global'). + embedding_model: Embedder model id. Defaults to ``EMBEDDING_MODEL`` + (GGUF nomic). The NPU profile passes the FLM-native embedder so + chat and embeddings stay co-resident on the NPU backend (#1744). + The embedding dimension is derived from the live embedder, not + this id, so a model with a different dim works without changes. Raises: RuntimeError: If Lemonade embedding service is unreachable @@ -331,6 +344,10 @@ def init_memory( # Explicit opt-out for environments that don't need memory (security # tests, lint-time imports, etc.). This is NOT a silent fallback — # the user/test author has explicitly set the env var. + self._embedding_model = embedding_model or EMBEDDING_MODEL + # Pre-probe default; refined from the live embedder below. + self._embedding_dim = EMBEDDING_DIM + if os.environ.get("GAIA_MEMORY_DISABLED") == "1": logger.info( "[MemoryMixin] memory disabled via GAIA_MEMORY_DISABLED=1; " @@ -405,17 +422,39 @@ def init_memory( # via ``_memory_store is None`` checks at every memory operation. try: self._get_embedder() - # Validate connectivity with a small test embedding + # Validate connectivity AND derive the embedding dimension from the + # live embedder — different embedders (e.g. the NPU FLM embedder) + # have different dims, so the FAISS index must match the active + # model rather than a hardcoded constant (#1744). test_vec = self._embed_text("connectivity test") - if test_vec.shape[0] != EMBEDDING_DIM: + dim = int(test_vec.shape[0]) + if dim <= 0: raise RuntimeError( - f"Embedding dimension mismatch: expected {EMBEDDING_DIM}, " - f"got {test_vec.shape[0]}" + f"Embedder '{self._embedding_model}' returned a 0-length " + "vector. Check that the model is loaded in Lemonade." ) + self._embedding_dim = dim logger.info( - "[MemoryMixin] Lemonade embedding service validated (%d-dim)", - EMBEDDING_DIM, + "[MemoryMixin] Lemonade embedding service validated " + "(model=%s, %d-dim)", + self._embedding_model, + self._embedding_dim, ) + # Invalidate stored vectors when the embedder changed. Vectors from + # a different model live in a different vector space (even at the + # same dim), so reusing them silently corrupts similarity search. + # Clearing forces backfill to re-embed with the active model. + prior = self._memory_store.get_embedder_id() + if prior is not None and prior != self._embedding_model: + cleared = self._memory_store.clear_all_embeddings() + logger.warning( + "[MemoryMixin] embedder changed (%s -> %s); cleared %d stored " + "embedding(s) for re-embedding", + prior, + self._embedding_model, + cleared, + ) + self._memory_store.set_embedder_id(self._embedding_model) except Exception as e: logger.warning( "[MemoryMixin] Lemonade embedding service unreachable — " @@ -633,7 +672,7 @@ def _get_embedder(self) -> Any: try: from gaia.llm.providers.lemonade import LemonadeProvider - self._embedder = LemonadeProvider(model=EMBEDDING_MODEL) + self._embedder = LemonadeProvider(model=self._embedding_model) logger.debug("[MemoryMixin] LemonadeProvider initialized for embeddings") return self._embedder except Exception as e: @@ -642,7 +681,7 @@ def _get_embedder(self) -> Any: ) from e def _embed_text(self, text: str) -> np.ndarray: - """Embed text via Lemonade (nomic-embed-text-v2-moe-GGUF, 768-dim). + """Embed text via Lemonade using the active embedder. Required, not optional. Raises RuntimeError if embedding fails. @@ -650,12 +689,12 @@ def _embed_text(self, text: str) -> np.ndarray: text: Text to embed. Returns: - L2-normalized float32 numpy array of shape (768,). + L2-normalized float32 numpy array of shape ``(self._embedding_dim,)``. """ embedder = self._get_embedder() try: # LemonadeProvider.embed() returns list[list[float]] - results = embedder.embed([text], model=EMBEDDING_MODEL) + results = embedder.embed([text], model=self._embedding_model) vec = np.array(results[0], dtype=np.float32) # L2-normalize for cosine similarity via IndexFlatIP @@ -713,13 +752,13 @@ def _rebuild_faiss_index(self) -> None: # Get all active knowledge items that have embeddings items = store.get_items_with_embeddings(include_sensitive=True) - index = faiss.IndexFlatIP(EMBEDDING_DIM) + index = faiss.IndexFlatIP(self._embedding_dim) id_map = [] for item in items: try: vec = _blob_to_embedding(item["embedding"]) - if vec.shape[0] != EMBEDDING_DIM: + if vec.shape[0] != self._embedding_dim: logger.debug( "[MemoryMixin] skipping embedding for %s: wrong dim %d", item["id"], @@ -779,11 +818,11 @@ def _faiss_remove(self, knowledge_id: str) -> None: # Reconstruct all vectors except the removed one n = self._faiss_index.ntotal if n <= 1: - self._faiss_index = faiss.IndexFlatIP(EMBEDDING_DIM) + self._faiss_index = faiss.IndexFlatIP(self._embedding_dim) self._faiss_id_map = [] return - all_vecs = np.zeros((n, EMBEDDING_DIM), dtype=np.float32) + all_vecs = np.zeros((n, self._embedding_dim), dtype=np.float32) for i in range(n): all_vecs[i] = self._faiss_index.reconstruct(i) @@ -791,7 +830,7 @@ def _faiss_remove(self, knowledge_id: str) -> None: keep_vecs = np.delete(all_vecs, idx, axis=0) keep_ids = self._faiss_id_map[:idx] + self._faiss_id_map[idx + 1 :] - new_index = faiss.IndexFlatIP(EMBEDDING_DIM) + new_index = faiss.IndexFlatIP(self._embedding_dim) new_index.add(keep_vecs) self._faiss_index = new_index self._faiss_id_map = keep_ids @@ -1478,7 +1517,7 @@ def reconcile_memory(self, max_pairs: int = 20) -> Dict: for item in items: try: vec = _blob_to_embedding(item["embedding"]) - if vec.shape[0] != EMBEDDING_DIM: + if vec.shape[0] != self._embedding_dim: continue norm = np.linalg.norm(vec) if norm > 0: @@ -1498,7 +1537,7 @@ def reconcile_memory(self, max_pairs: int = 20) -> Dict: try: import faiss - temp_index = faiss.IndexFlatIP(EMBEDDING_DIM) + temp_index = faiss.IndexFlatIP(self._embedding_dim) temp_index.add(mat) # Search each item for its top-5 neighbors diff --git a/src/gaia/agents/base/memory_store.py b/src/gaia/agents/base/memory_store.py index 6dfb17110..282d5a564 100644 --- a/src/gaia/agents/base/memory_store.py +++ b/src/gaia/agents/base/memory_store.py @@ -248,6 +248,15 @@ def _safe_json_loads(value) -> object: -- Knowledge FTS5 (standalone, manually synced, porter stemmer for morphological matching) CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_fts USING fts5(content, domain, category, tokenize='porter unicode61'); +-- Key/value metadata (e.g. the embedder id that produced stored vectors). +-- Lives in the DB (not a sidecar file) so it is atomic with the data it +-- describes and visible to every connection — the agent and the UI memory +-- router open the same DB from different code paths (#1744). +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT +); + -- Tool history CREATE TABLE IF NOT EXISTS tool_history ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -1460,6 +1469,55 @@ def store_embedding(self, knowledge_id: str, embedding: bytes) -> bool: self._conn.rollback() raise + def clear_all_embeddings(self) -> int: + """Null out every stored embedding so they get re-embedded on backfill. + + Used when the active embedder changes: vectors from a different model + live in a different vector space (and possibly a different dimension), + so reusing them would silently corrupt similarity search. Returns the + number of rows cleared. + """ + with self._lock: + try: + rowcount = self._conn.execute( + "UPDATE knowledge SET embedding = NULL WHERE embedding IS NOT NULL" + ).rowcount + self._conn.commit() + return rowcount + except Exception: + self._conn.rollback() + raise + + #: ``meta`` key recording which embedder produced the stored vectors. + _EMBEDDER_META_KEY = "embedder_id" + + def get_embedder_id(self) -> Optional[str]: + """Return the embedder model id that produced the stored embeddings. + + ``None`` when nothing has been stamped yet (fresh DB) — callers treat + that as "no change to detect". Read from the ``meta`` table so every + connection (agent + UI router) sees the same value. + """ + with self._lock: + row = self._conn.execute( + "SELECT value FROM meta WHERE key = ?", (self._EMBEDDER_META_KEY,) + ).fetchone() + return row[0] if row else None + + def set_embedder_id(self, model_id: str) -> None: + """Record the embedder model id that produced the stored embeddings.""" + with self._lock: + try: + self._conn.execute( + "INSERT INTO meta (key, value) VALUES (?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (self._EMBEDDER_META_KEY, model_id), + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise + def get_items_with_embeddings( self, category: str | None = None, diff --git a/src/gaia/agents/registry.py b/src/gaia/agents/registry.py index 155a95569..c1c97c679 100644 --- a/src/gaia/agents/registry.py +++ b/src/gaia/agents/registry.py @@ -217,6 +217,11 @@ def _compute_custom_origin_hash(py_file: Path) -> str: return hashlib.sha256(py_file.read_bytes()).hexdigest()[:16] +# Default embedder (GPU/CPU). The NPU device overrides this with the FLM-native +# embedder so chat + embeddings stay co-resident on the NPU backend (#1744). +DEFAULT_EMBEDDING_MODEL = "nomic-embed-text-v2-moe-GGUF" + + @dataclass class DeviceConfig: """A verified (device, model, recipe, backend) configuration for an agent. @@ -235,6 +240,10 @@ class DeviceConfig: verified: Whether this combination has been tested end-to-end via agent eval. Unverified configs show a warning badge in the UI. ctx_size: Default context window size for this configuration. + embedding_model: Embedder model id for RAG/memory on this device. NPU + uses the FLM-native embedder so the chat model and embedder stay + co-resident on the NPU backend; a GGUF embedder runs on Vulkan and + evicts the FLM chat model every turn on a shared-memory APU (#1744). """ device: Literal["cpu", "gpu", "npu"] @@ -243,6 +252,7 @@ class DeviceConfig: backend: str verified: bool = False ctx_size: int = 32768 + embedding_model: str = DEFAULT_EMBEDDING_MODEL # Default device configurations for built-in agents using Gemma 4 E4B. @@ -271,10 +281,28 @@ class DeviceConfig: backend="flm:npu", verified=True, ctx_size=32768, + # FLM-native embedder so chat + embeddings stay co-resident on the NPU + # backend and don't thrash NPU<->Vulkan every turn (#1744). + embedding_model="embed-gemma-300m-FLM", ), ] +def get_embedding_model_for_device(device: Optional[str]) -> str: + """Return the embedder model id for a device target. + + Single source of truth: reads ``DEFAULT_DEVICE_CONFIGS`` so the embedder + choice lives next to the chat model/recipe/backend for each device. The NPU + profile uses the FLM-native embedder (see ``DeviceConfig.embedding_model``); + GPU/CPU and an unspecified device default to the GGUF nomic embedder, which + matches the GPU-default policy elsewhere in the CLI. + """ + for dc in DEFAULT_DEVICE_CONFIGS: + if dc.device == device: + return dc.embedding_model + return DEFAULT_EMBEDDING_MODEL + + @dataclass class ModelTier: """A selectable model size for an agent (#1162). diff --git a/src/gaia/installer/init_command.py b/src/gaia/installer/init_command.py index 55456f034..bc7dac90f 100644 --- a/src/gaia/installer/init_command.py +++ b/src/gaia/installer/init_command.py @@ -98,7 +98,11 @@ "npu": { "description": "Ryzen AI NPU acceleration via FLM backend (requires XDNA2 NPU)", "agent": "chat", - "models": ["gemma4-it-e2b-FLM"], + # FLM chat model + FLM-native embedder so chat and embeddings stay + # co-resident on the NPU backend. A GGUF embedder would run on Vulkan + # and evict the FLM chat model every turn (#1744). Both are built-in + # Lemonade *-FLM models, pulled by name only (no recipe — #1655). + "models": ["gemma4-it-e2b-FLM", "embed-gemma-300m-FLM"], "approx_size": "~3 GB", "min_lemonade_version": "10.2.0", # NPU context window. Matches GPU/CPU (32768) so the init report and diff --git a/src/gaia/llm/lemonade_client.py b/src/gaia/llm/lemonade_client.py index 5a41a3b2a..ce3909a50 100644 --- a/src/gaia/llm/lemonade_client.py +++ b/src/gaia/llm/lemonade_client.py @@ -283,6 +283,21 @@ class LemonadeStatus: min_ctx_size=2048, tool_calling=False, ), + # --- NPU-native FLM embedder for the NPU profile (#1744) --- + # EmbeddingGemma 300M built for the FastFlowLM/NPU backend. On a shared- + # memory Ryzen AI APU the GGUF nomic embedder runs on Vulkan/llama.cpp and + # reclaims the memory the FLM chat model holds, so loading it evicts the + # chat model — every chat turn then thrashes NPU<->Vulkan (#1676). Keeping + # the embedder on the same FLM/NPU backend as the chat model lets both stay + # co-resident. Built-in Lemonade *-FLM model: pull by name only (no recipe; + # passing recipe triggers user-model registration and 400s — #1655). + "embed-gemma-flm": ModelRequirement( + model_type=ModelType.EMBEDDING, + model_id="embed-gemma-300m-FLM", + display_name="EmbeddingGemma 300M (NPU/FLM)", + min_ctx_size=2048, + tool_calling=False, + ), } # Define agent profiles with their model requirements diff --git a/src/gaia/ui/routers/memory.py b/src/gaia/ui/routers/memory.py index a35fc1517..203f0e1d2 100644 --- a/src/gaia/ui/routers/memory.py +++ b/src/gaia/ui/routers/memory.py @@ -691,17 +691,23 @@ def rebuild_embeddings() -> Dict: from gaia.agents.base.memory import EMBEDDING_MODEL from gaia.llm.providers.lemonade import LemonadeProvider - provider = LemonadeProvider(model=EMBEDDING_MODEL) + store = _get_store() + # Re-embed with the embedder that produced the stored vectors, not a + # hardcoded default. On the NPU profile that is the FLM-native embedder; + # using nomic here would mix vector spaces in one table and reload a + # Vulkan GGUF embedder that evicts the FLM chat model (#1744). + embedder_model = store.get_embedder_id() or EMBEDDING_MODEL + provider = LemonadeProvider(model=embedder_model) def _embed_fn(text: str) -> bytes: - results = provider.embed([text], model=EMBEDDING_MODEL) + results = provider.embed([text], model=embedder_model) vec = np.array(results[0], dtype=np.float32) norm = np.linalg.norm(vec) if norm > 0: vec = vec / norm return vec.astype(np.float32).tobytes() - result = _get_store().backfill_embeddings(_embed_fn) + result = store.backfill_embeddings(_embed_fn) return result except Exception as exc: logger.error("[memory router] rebuild-embeddings failed: %s", exc) @@ -757,7 +763,7 @@ def trigger_reconciliation(max_pairs: int = Query(20, ge=1, le=100)) -> Dict: except ImportError: raise HTTPException(503, "faiss not installed — run: pip install faiss-cpu") - from gaia.agents.base.memory import EMBEDDING_DIM, _blob_to_embedding + from gaia.agents.base.memory import _blob_to_embedding from gaia.llm import create_client store = _get_store() @@ -766,15 +772,21 @@ def trigger_reconciliation(max_pairs: int = Query(20, ge=1, le=100)) -> Dict: if len(items) < 2: return result # Not enough items with embeddings - # Build id list and vector matrix + # Dimension is derived from the stored vectors, not a hardcoded constant + # — the active embedder (e.g. the NPU FLM embedder) may differ from the + # nomic default (#1744). The first valid vector sets the expected dim; + # any vector of a different dim (a stale pre-switch embedding) is skipped. ids: List[str] = [] vectors: List[np.ndarray] = [] item_map: Dict[str, Any] = {} + expected_dim: Optional[int] = None for item in items: try: vec = _blob_to_embedding(item["embedding"]) - if vec.shape[0] != EMBEDDING_DIM: + if expected_dim is None: + expected_dim = int(vec.shape[0]) + if vec.shape[0] != expected_dim: continue norm = np.linalg.norm(vec) if norm > 0: @@ -785,11 +797,11 @@ def trigger_reconciliation(max_pairs: int = Query(20, ge=1, le=100)) -> Dict: except Exception: continue - if len(vectors) < 2: + if len(vectors) < 2 or expected_dim is None: return result mat = np.stack(vectors).astype(np.float32) - index = faiss.IndexFlatIP(EMBEDDING_DIM) + index = faiss.IndexFlatIP(expected_dim) index.add(mat) n_neighbors = min(5, len(vectors)) diff --git a/tests/test_gemma_embedder_hw.py b/tests/test_gemma_embedder_hw.py new file mode 100644 index 000000000..48697716c --- /dev/null +++ b/tests/test_gemma_embedder_hw.py @@ -0,0 +1,164 @@ +# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +"""Hardware validation for the NPU-native FLM embedder (#1744). + +Unlike ``tests/unit/test_npu_flm_embedder.py`` (fully mocked), these tests run +``embed-gemma-300m-FLM`` on a **real** FLM/NPU-enabled Lemonade Server and prove +the feature's actual purpose: the FLM embedder produces valid vectors on the NPU +and stays co-resident with the FLM chat model (no NPU<->Vulkan eviction). + +Gating: ``embed-gemma-300m-FLM`` is an FLM-recipe model that only exists on a +Lemonade Server running the FLM backend on Ryzen AI NPU hardware -- it is not in +the mainline llamacpp/Vulkan catalog. Every test here therefore SKIPS unless the +live server actually advertises the model, so the suite is safe on CPU/GPU boxes +and in Vulkan CI while still executing for real on an NPU runner. + +Run on NPU hardware with: + python -m pytest tests/test_npu_flm_embedder_hw.py -v -rs +""" + +import math + +import pytest + +from gaia.agents.registry import get_embedding_model_for_device +from gaia.llm.lemonade_client import LemonadeClient, LemonadeClientError + +pytestmark = [pytest.mark.integration, pytest.mark.real_model] + +# The NPU profile's embedder id, resolved from the same source the agent uses so +# a rename in the registry can't silently desync this test from production. +FLM_EMBEDDER = get_embedding_model_for_device("npu") # "embed-gemma-300m-FLM" +FLM_CHAT_MODEL = "gemma4-it-e2b-FLM" # NPU profile chat model (registry.py) + +# EmbeddingGemma 300M is 768-dim; the agent derives this at runtime rather than +# hardcoding it, so a mismatch here means the wrong model/vector space loaded. +EXPECTED_DIM = 768 + + +def _catalog_ids(client: LemonadeClient) -> set: + """Model ids the live server advertises (empty set if unreachable).""" + try: + catalog = client.list_models(show_all=True) + except LemonadeClientError: + return set() + data = catalog.get("data", catalog) if isinstance(catalog, dict) else catalog + ids = set() + for entry in data or []: + if isinstance(entry, dict) and entry.get("id"): + ids.add(entry["id"]) + return ids + + +@pytest.fixture(scope="module") +def client(): + return LemonadeClient() + + +@pytest.fixture +def npu_embedder(client, require_lemonade): + """Skip unless the FLM embedder is available on the live server, then load it. + + ``require_lemonade`` skips when no server is up; this then skips when the + server is up but not FLM/NPU-enabled (the Vulkan-catalog case). + """ + if FLM_EMBEDDER not in _catalog_ids(client): + pytest.skip( + f"{FLM_EMBEDDER} not in the live Lemonade catalog -- requires an " + "FLM/NPU-enabled server on Ryzen AI hardware" + ) + # Built-in FLM model: pull/load by name, no recipe (#1655-safe). + client.load_model(FLM_EMBEDDER, auto_download=True, prompt=False) + return FLM_EMBEDDER + + +class TestFlmEmbedderOnNpu: + def test_single_embedding_is_valid(self, client, npu_embedder): + """One text -> one finite, non-zero, correctly-sized vector.""" + resp = client.embeddings( + "resilient local AI on the NPU", model=npu_embedder, timeout=120 + ) + + assert isinstance(resp, dict) and "data" in resp + assert len(resp["data"]) == 1 + vec = resp["data"][0]["embedding"] + assert len(vec) == EXPECTED_DIM, f"expected {EXPECTED_DIM}-dim, got {len(vec)}" + assert all(isinstance(x, float) and math.isfinite(x) for x in vec) + assert any(x != 0.0 for x in vec), "embedding must not be all zeros" + + def test_batch_embedding_shapes(self, client, npu_embedder): + """A batch returns one same-dim vector per input, in order.""" + texts = [ + "The NPU runs the FLM embedder.", + "Vulkan runs the GGUF embedder.", + "Chat and embeddings stay co-resident.", + ] + resp = client.embeddings(texts, model=npu_embedder, timeout=180) + + assert len(resp["data"]) == len(texts) + dims = {len(item["embedding"]) for item in resp["data"]} + assert dims == {EXPECTED_DIM}, f"inconsistent/unexpected dims: {dims}" + + def test_embedding_is_deterministic(self, client, npu_embedder): + """Same input -> identical vector (no sampling in embeddings).""" + text = "determinism check" + a = client.embeddings([text], model=npu_embedder, timeout=120)["data"][0][ + "embedding" + ] + b = client.embeddings([text], model=npu_embedder, timeout=120)["data"][0][ + "embedding" + ] + assert a == b + + def test_distinct_texts_differ(self, client, npu_embedder): + """Different inputs -> different vectors (model is actually encoding).""" + a = client.embeddings(["apples"], model=npu_embedder, timeout=120)["data"][0][ + "embedding" + ] + b = client.embeddings( + ["quantum chromodynamics"], model=npu_embedder, timeout=120 + )["data"][0]["embedding"] + assert a != b + + +class TestFlmCoresidency: + """The core #1744 guarantee: the FLM chat model and FLM embedder are both + resident at once, so a chat turn does not evict the embedder (and vice + versa) the way the Vulkan GGUF embedder did on a shared-memory APU.""" + + def test_chat_and_embedder_coresident(self, client, npu_embedder): + if FLM_CHAT_MODEL not in _catalog_ids(client): + pytest.skip( + f"{FLM_CHAT_MODEL} not in the live catalog -- cannot test co-residency" + ) + + client.load_model( + FLM_CHAT_MODEL, auto_download=True, prompt=False, ctx_size=32768 + ) + # Touch the embedder after the chat model loads; on the old Vulkan path + # this is exactly the call that evicted the FLM chat model. + client.load_model(npu_embedder, auto_download=True, prompt=False) + + health = client._send_request("GET", f"{client.base_url}/health", timeout=15) + loaded = health.get("all_models_loaded") or [] + loaded_ids = { + ( + (m.get("id") or m.get("model_name") or m.get("model")) + if isinstance(m, dict) + else m + ) + for m in loaded + } + + assert FLM_EMBEDDER in loaded_ids, ( + f"{FLM_EMBEDDER} not resident after loading the chat model " + f"(loaded={sorted(str(i) for i in loaded_ids)}) -- NPU<->Vulkan eviction regressed" + ) + assert FLM_CHAT_MODEL in loaded_ids, ( + f"{FLM_CHAT_MODEL} was evicted when the embedder loaded " + f"(loaded={sorted(str(i) for i in loaded_ids)})" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s", "-rs"]) diff --git a/tests/unit/test_init_command.py b/tests/unit/test_init_command.py index 45a7fbbcb..b0b2c7729 100644 --- a/tests/unit/test_init_command.py +++ b/tests/unit/test_init_command.py @@ -420,9 +420,10 @@ def test_force_models_deletes_before_download(self, mock_installer_class): def test_npu_profile_pulls_builtin_model_without_recipe(self, mock_installer_class): """NPU/FLM models are built-in; pulling with a recipe 400s (#1655). - The npu profile must download ``gemma4-it-e2b-FLM`` via - ensure_model_downloaded (pull by name), never pull_model(recipe=...), - which Lemonade rejects unless the name carries a ``user.`` prefix. + The npu profile must download both the FLM chat model and the FLM-native + embedder (#1744) via ensure_model_downloaded (pull by name), never + pull_model(recipe=...), which Lemonade rejects unless the name carries a + ``user.`` prefix. """ from gaia.installer.init_command import InitCommand @@ -435,9 +436,10 @@ def test_npu_profile_pulls_builtin_model_without_recipe(self, mock_installer_cla result = cmd._download_models() self.assertTrue(result) - mock_client.ensure_model_downloaded.assert_called_once_with( - "gemma4-it-e2b-FLM" - ) + pulled = { + c.args[0] for c in mock_client.ensure_model_downloaded.call_args_list + } + self.assertEqual(pulled, {"gemma4-it-e2b-FLM", "embed-gemma-300m-FLM"}) # Regression guard: no recipe-bearing pull_model call. mock_client.pull_model.assert_not_called() diff --git a/tests/unit/test_memory_router.py b/tests/unit/test_memory_router.py index 2e28d9945..005224936 100644 --- a/tests/unit/test_memory_router.py +++ b/tests/unit/test_memory_router.py @@ -820,6 +820,35 @@ def boom(*args, **kwargs): resp = client.post("/api/memory/rebuild-embeddings") assert resp.status_code == 500 + def test_rebuild_uses_stamped_embedder_not_nomic(self, client, test_store): + """Re-embed with the stamped embedder, not the nomic default (#1744). + + On the NPU profile the agent stamps the FLM embedder. The dashboard + rebuild button must reuse it — embedding with nomic would mix vector + spaces in one table and reload a Vulkan embedder that evicts the FLM + chat model. + """ + test_store.set_embedder_id("embed-gemma-300m-FLM") + mock_provider = MagicMock() + mock_provider.embed.return_value = [[0.1] * 768] + with patch( + "gaia.llm.providers.lemonade.LemonadeProvider", return_value=mock_provider + ) as ctor: + resp = client.post("/api/memory/rebuild-embeddings") + assert resp.status_code == 200 + ctor.assert_called_once_with(model="embed-gemma-300m-FLM") + + def test_rebuild_falls_back_to_nomic_when_unstamped(self, client, test_store): + """A never-stamped DB (no agent run yet) defaults to the nomic embedder.""" + mock_provider = MagicMock() + mock_provider.embed.return_value = [[0.1] * 768] + with patch( + "gaia.llm.providers.lemonade.LemonadeProvider", return_value=mock_provider + ) as ctor: + resp = client.post("/api/memory/rebuild-embeddings") + assert resp.status_code == 200 + ctor.assert_called_once_with(model="nomic-embed-text-v2-moe-GGUF") + # =========================================================================== # v2 Router Tests — Reconcile Endpoint @@ -842,6 +871,31 @@ def test_reconcile_returns_503_when_no_agent_and_faiss_missing(self, client): assert resp.status_code == 503 assert "faiss" in resp.json()["detail"].lower() + def test_reconcile_standalone_derives_dim_from_vectors(self, client, test_store): + """Standalone reconcile builds the FAISS index at the stored vectors' + dim, not a hardcoded 768 — so a non-768 embedder (e.g. a truncated FLM + embedder) is not silently skipped (#1744). + """ + pytest.importorskip("faiss") # standalone reconcile path needs faiss + import numpy as np + + # Two identical 512-dim vectors → cosine 1.0 → above the pair threshold. + vec = np.ones(512, dtype=np.float32).tobytes() + for content in ("alpha fact one", "alpha fact two"): + kid = _create_knowledge(client, content)["knowledge_id"] + test_store.store_embedding(kid, vec) + test_store.set_embedder_id("embed-gemma-300m-FLM") + + fake_llm = MagicMock() + fake_llm.chat.return_value = '{"relationship": "neutral", "action": "noop"}' + with patch("gaia.llm.create_client", return_value=fake_llm): + resp = client.post("/api/memory/reconcile") + + assert resp.status_code == 200 + # A pair was found and classified — proves the 512-dim index built and + # nothing was skipped by a dim mismatch. + assert resp.json()["pairs_checked"] >= 1 + def test_reconcile_returns_200_when_agent_registered(self, client): """POST /api/memory/reconcile returns 200 when _reconcile_fn is set.""" memory_router_mod._reconcile_fn = lambda: { diff --git a/tests/unit/test_npu_flm_embedder.py b/tests/unit/test_npu_flm_embedder.py new file mode 100644 index 000000000..f55f11597 --- /dev/null +++ b/tests/unit/test_npu_flm_embedder.py @@ -0,0 +1,249 @@ +# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +"""Unit tests for the NPU-native FLM embedder wiring (#1744). + +Covers: +- The device->embedder resolver and DeviceConfig.embedding_model field. +- The embed-gemma-300m-FLM entry in the Lemonade MODELS registry. +- The NPU init profile downloading the FLM embedder. +- MemoryMixin deriving the embedding dimension from the live embedder + (no hardcoded 768) and invalidating stored vectors when the embedder + model changes (same dim, different vector space). + +All tests mock the embedder; no Lemonade server or NPU required. +""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from gaia.agents.base.memory_store import MemoryStore + +# --------------------------------------------------------------------------- +# Device -> embedder resolution +# --------------------------------------------------------------------------- + + +class TestDeviceEmbedderResolution: + def test_npu_uses_flm_embedder(self): + from gaia.agents.registry import get_embedding_model_for_device + + assert get_embedding_model_for_device("npu") == "embed-gemma-300m-FLM" + + @pytest.mark.parametrize("device", ["gpu", "cpu", None, "unknown-device"]) + def test_non_npu_uses_nomic(self, device): + from gaia.agents.registry import get_embedding_model_for_device + + assert get_embedding_model_for_device(device) == "nomic-embed-text-v2-moe-GGUF" + + def test_device_config_carries_embedder(self): + """The embedder choice lives next to the chat model in DeviceConfig.""" + from gaia.agents.registry import DEFAULT_DEVICE_CONFIGS + + by_device = {dc.device: dc for dc in DEFAULT_DEVICE_CONFIGS} + assert by_device["npu"].embedding_model == "embed-gemma-300m-FLM" + assert by_device["gpu"].embedding_model == "nomic-embed-text-v2-moe-GGUF" + assert by_device["cpu"].embedding_model == "nomic-embed-text-v2-moe-GGUF" + + +# --------------------------------------------------------------------------- +# Lemonade model registry + download manifest +# --------------------------------------------------------------------------- + + +class TestModelRegistry: + def test_flm_embedder_registered(self): + from gaia.llm.lemonade_client import MODELS, ModelType + + mr = MODELS["embed-gemma-flm"] + assert mr.model_id == "embed-gemma-300m-FLM" + assert mr.model_type == ModelType.EMBEDDING + # FLM/NPU server 500s on a tools payload, like the e2b chat model. + assert mr.tool_calling is False + + def test_npu_init_profile_downloads_flm_embedder(self): + from gaia.installer.init_command import INIT_PROFILES + + models = INIT_PROFILES["npu"]["models"] + assert "embed-gemma-300m-FLM" in models + # Chat model must remain so both pull on `gaia init --profile npu`. + assert "gemma4-it-e2b-FLM" in models + + +# --------------------------------------------------------------------------- +# MemoryMixin: dynamic dim + embedder-change invalidation +# --------------------------------------------------------------------------- + + +class _FakeAgent: + """Minimal host so MemoryMixin can be exercised without a real Agent.""" + + def __init__(self): + self._system_prompt_cache = None + + def register_tool(self, *a, **k): + pass + + +def _make_host(): + from gaia.agents.base.memory import MemoryMixin + + class _Host(MemoryMixin, _FakeAgent): + pass + + return _Host() + + +def _init_with_embedder(host, db_path, *, model, dim): + """Run init_memory with a mock embedder of a given model id and dim.""" + from gaia.agents.base.memory import MemoryMixin + + vec = np.random.rand(dim).astype(np.float32) + with ( + patch.object(MemoryMixin, "_get_embedder", return_value=MagicMock()), + patch.object(MemoryMixin, "_embed_text", return_value=vec), + patch.object(MemoryMixin, "_backfill_embeddings", return_value=0), + patch.object(MemoryMixin, "init_system_context", return_value=None), + ): + host.init_memory(db_path=db_path, embedding_model=model) + + +class TestDynamicDim: + def test_dim_derived_from_live_embedder(self, tmp_path, monkeypatch): + """A 512-dim embedder yields a 512-dim FAISS index — no hardcoded 768.""" + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + host = _make_host() + _init_with_embedder( + host, tmp_path / "memory.db", model="embed-gemma-300m-FLM", dim=512 + ) + assert host._embedding_dim == 512 + assert host._embedding_model == "embed-gemma-300m-FLM" + # When faiss is available, the index is built at the derived dim, not + # the module default. faiss is optional, so skip this leg without it. + pytest.importorskip("faiss") + assert host._faiss_index is not None + assert host._faiss_index.d == 512 + + def test_default_embedder_when_unspecified(self, tmp_path, monkeypatch): + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + host = _make_host() + _init_with_embedder(host, tmp_path / "memory.db", model=None, dim=768) + assert host._embedding_model == "nomic-embed-text-v2-moe-GGUF" + assert host._embedding_dim == 768 + + def test_zero_length_vector_degrades_memory(self, tmp_path, monkeypatch): + """A 0-length embedding trips the guard and degrades memory loudly, + rather than building a malformed index.""" + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + from gaia.agents.base.memory import MemoryMixin + + host = _make_host() + empty = np.empty(0, dtype=np.float32) + with ( + patch.object(MemoryMixin, "_get_embedder", return_value=MagicMock()), + patch.object(MemoryMixin, "_embed_text", return_value=empty), + patch.object(MemoryMixin, "init_system_context", return_value=None), + ): + host.init_memory( + db_path=tmp_path / "memory.db", embedding_model="broken-embedder" + ) + # The RuntimeError is caught and memory is disabled for the session. + assert host._memory_store is None + + +class TestEmbedderChangeInvalidation: + def test_switching_embedder_clears_stored_vectors(self, tmp_path, monkeypatch): + """Same dim, different model => stored vectors must be invalidated.""" + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + db_path = tmp_path / "memory.db" + + # Seed a store as if a prior nomic session wrote a 768-dim vector. + store = MemoryStore(db_path=db_path) + kid = store.store(category="fact", content="the sky is blue") + store.store_embedding(kid, np.zeros(768, dtype=np.float32).tobytes()) + store.set_embedder_id("nomic-embed-text-v2-moe-GGUF") + assert store.get_embedding_coverage()["without_embedding"] == 0 + store.close() + + # New session on the FLM embedder (also 768-dim, different space). + host = _make_host() + _init_with_embedder(host, db_path, model="embed-gemma-300m-FLM", dim=768) + + # Marker updated and the stale vector was cleared for re-embedding. + assert host._memory_store.get_embedder_id() == "embed-gemma-300m-FLM" + assert host._memory_store.get_embedding_coverage()["without_embedding"] == 1 + + def test_switch_then_backfill_repopulates(self, tmp_path, monkeypatch): + """After a switch the cleared vectors are re-embedded with the new model + (full round trip: clear -> backfill -> coverage restored).""" + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + from gaia.agents.base.memory import MemoryMixin + + db_path = tmp_path / "memory.db" + store = MemoryStore(db_path=db_path) + kid = store.store(category="fact", content="the sky is blue") + store.store_embedding(kid, np.zeros(768, dtype=np.float32).tobytes()) + store.set_embedder_id("nomic-embed-text-v2-moe-GGUF") + store.close() + + # Real backfill this time (only _embed_text is mocked). + host = _make_host() + vec = np.random.rand(768).astype(np.float32) + with ( + patch.object(MemoryMixin, "_get_embedder", return_value=MagicMock()), + patch.object(MemoryMixin, "_embed_text", return_value=vec), + patch.object(MemoryMixin, "init_system_context", return_value=None), + ): + host.init_memory(db_path=db_path, embedding_model="embed-gemma-300m-FLM") + + cov = host._memory_store.get_embedding_coverage() + assert cov["without_embedding"] == 0 # re-embedded with the new model + assert host._memory_store.get_embedder_id() == "embed-gemma-300m-FLM" + + def test_same_embedder_keeps_vectors(self, tmp_path, monkeypatch): + monkeypatch.delenv("GAIA_MEMORY_DISABLED", raising=False) + db_path = tmp_path / "memory.db" + + store = MemoryStore(db_path=db_path) + kid = store.store(category="fact", content="water is wet") + store.store_embedding(kid, np.zeros(768, dtype=np.float32).tobytes()) + store.set_embedder_id("nomic-embed-text-v2-moe-GGUF") + store.close() + + host = _make_host() + with patch( + "gaia.agents.base.memory_store.MemoryStore.clear_all_embeddings" + ) as clear: + _init_with_embedder( + host, db_path, model="nomic-embed-text-v2-moe-GGUF", dim=768 + ) + clear.assert_not_called() + + +# --------------------------------------------------------------------------- +# MemoryStore embedder marker + clear +# --------------------------------------------------------------------------- + + +class TestMemoryStoreEmbedderMarker: + def test_marker_round_trip(self, tmp_path): + store = MemoryStore(db_path=tmp_path / "memory.db") + try: + assert store.get_embedder_id() is None # fresh DB + store.set_embedder_id("embed-gemma-300m-FLM") + assert store.get_embedder_id() == "embed-gemma-300m-FLM" + finally: + store.close() + + def test_clear_all_embeddings(self, tmp_path): + store = MemoryStore(db_path=tmp_path / "memory.db") + try: + kid = store.store(category="fact", content="hello") + store.store_embedding(kid, np.zeros(768, dtype=np.float32).tobytes()) + assert store.get_embedding_coverage()["without_embedding"] == 0 + cleared = store.clear_all_embeddings() + assert cleared == 1 + assert store.get_embedding_coverage()["without_embedding"] == 1 + finally: + store.close()