csabhay commited on
Commit
4e3ffdb
·
1 Parent(s): c25cdee

fix: track active embed provider to prevent dimension mismatch on query

Browse files
Files changed (1) hide show
  1. 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 a batch of texts using the configured provider."""
69
- if self.embed_provider == "gemini":
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), falling back to HF", e)
 
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}