Spaces:
Sleeping
Sleeping
fix: track active embed provider to prevent dimension mismatch on query
Browse files- backend/rag.py +6 -3
backend/rag.py
CHANGED
|
@@ -50,6 +50,7 @@ class RAGEngine:
|
|
| 50 |
self.top_k = top_k
|
| 51 |
self.llm_model = llm_model
|
| 52 |
self.embed_provider = embed_provider
|
|
|
|
| 53 |
self.chunks: list[str] = []
|
| 54 |
self.chunk_embeddings: np.ndarray | None = None
|
| 55 |
self._ingest_progress: dict = {"state": "idle", "embedded": 0, "total": 0}
|
|
@@ -65,8 +66,8 @@ class RAGEngine:
|
|
| 65 |
self._genai_client = None
|
| 66 |
|
| 67 |
def _embed_texts(self, texts: list[str], task_type: str = "RETRIEVAL_DOCUMENT") -> np.ndarray:
|
| 68 |
-
"""Embed
|
| 69 |
-
if self.
|
| 70 |
return self._embed_gemini(texts, task_type)
|
| 71 |
return self._embed_hf(texts)
|
| 72 |
|
|
@@ -82,7 +83,8 @@ class RAGEngine:
|
|
| 82 |
emb = emb.reshape(1, -1)
|
| 83 |
return emb
|
| 84 |
except Exception as e:
|
| 85 |
-
log.warning("Gemini embedding failed (%s),
|
|
|
|
| 86 |
return self._embed_hf(texts)
|
| 87 |
|
| 88 |
def _embed_hf(self, texts: list[str]) -> np.ndarray:
|
|
@@ -128,6 +130,7 @@ class RAGEngine:
|
|
| 128 |
)
|
| 129 |
self.chunks = splitter.split_text(text)
|
| 130 |
self.chunk_embeddings = None
|
|
|
|
| 131 |
self._ingest_generation += 1
|
| 132 |
total = len(self.chunks)
|
| 133 |
self._ingest_progress = {"state": "running", "embedded": 0, "total": total}
|
|
|
|
| 50 |
self.top_k = top_k
|
| 51 |
self.llm_model = llm_model
|
| 52 |
self.embed_provider = embed_provider
|
| 53 |
+
self._active_embed_provider = embed_provider # tracks what was actually used
|
| 54 |
self.chunks: list[str] = []
|
| 55 |
self.chunk_embeddings: np.ndarray | None = None
|
| 56 |
self._ingest_progress: dict = {"state": "idle", "embedded": 0, "total": 0}
|
|
|
|
| 66 |
self._genai_client = None
|
| 67 |
|
| 68 |
def _embed_texts(self, texts: list[str], task_type: str = "RETRIEVAL_DOCUMENT") -> np.ndarray:
|
| 69 |
+
"""Embed using the active provider (locked in after first successful batch)."""
|
| 70 |
+
if self._active_embed_provider == "gemini":
|
| 71 |
return self._embed_gemini(texts, task_type)
|
| 72 |
return self._embed_hf(texts)
|
| 73 |
|
|
|
|
| 83 |
emb = emb.reshape(1, -1)
|
| 84 |
return emb
|
| 85 |
except Exception as e:
|
| 86 |
+
log.warning("Gemini embedding failed (%s), switching to HF for this session", e)
|
| 87 |
+
self._active_embed_provider = "hf" # lock to HF for all remaining batches + queries
|
| 88 |
return self._embed_hf(texts)
|
| 89 |
|
| 90 |
def _embed_hf(self, texts: list[str]) -> np.ndarray:
|
|
|
|
| 130 |
)
|
| 131 |
self.chunks = splitter.split_text(text)
|
| 132 |
self.chunk_embeddings = None
|
| 133 |
+
self._active_embed_provider = self.embed_provider # reset for each fresh ingestion
|
| 134 |
self._ingest_generation += 1
|
| 135 |
total = len(self.chunks)
|
| 136 |
self._ingest_progress = {"state": "running", "embedded": 0, "total": total}
|