@@ -135,6 +135,16 @@ async def add_documents(
135135 if not uncached :
136136 return 0
137137
138+ # Cost-safety gate. Embedding a corpus is cheap-per-doc but can
139+ # add up over thousands of files; the rate limit also exposes a
140+ # real-money path through misconfigured loops. Gate per-document.
141+ try :
142+ from deepr .experts .cost_safety import get_cost_safety_manager
143+
144+ _cost_safety = get_cost_safety_manager ()
145+ except Exception :
146+ _cost_safety = None # type: ignore[assignment]
147+
138148 # Batch embed new documents (more efficient than one-by-one)
139149 new_embeddings = []
140150 new_metadata = []
@@ -146,9 +156,38 @@ async def add_documents(
146156 # Truncate content for embedding (model limit)
147157 embed_content = content [:8000 ]
148158
159+ if _cost_safety is not None :
160+ try :
161+ _est = 0.0002 # text-embedding-3-small ~$0.02/M tokens
162+ _allowed , _reason , _ = _cost_safety .check_operation (
163+ session_id = f"embed:{ self .expert_name } " ,
164+ operation_type = "embed_document" ,
165+ estimated_cost = _est ,
166+ require_confirmation = False ,
167+ )
168+ if not _allowed :
169+ logger .warning ("Embedding for %s blocked by cost-safety: %s" , filename , _reason )
170+ continue
171+ except Exception :
172+ _est = 0.0
173+ else :
174+ _est = 0.0
175+
149176 try :
150177 response = await client .embeddings .create (model = model , input = embed_content )
151178 embedding = np .array (response .data [0 ].embedding )
179+ if _cost_safety is not None :
180+ try :
181+ _cost_safety .record_cost (
182+ session_id = f"embed:{ self .expert_name } " ,
183+ operation_type = "embed_document" ,
184+ actual_cost = float (_est ),
185+ provider = "openai" ,
186+ model = model ,
187+ source = "experts.embedding_cache.add_documents" ,
188+ )
189+ except Exception :
190+ pass
152191
153192 content_hash = self ._content_hash (content )
154193
@@ -205,10 +244,41 @@ async def search(self, query: str, client, top_k: int = 5, model: str = "text-em
205244 if self .embeddings is None or len (self .embeddings ) == 0 :
206245 return []
207246
247+ # Cost-safety gate for query embedding.
248+ try :
249+ from deepr .experts .cost_safety import get_cost_safety_manager
250+
251+ _cost_safety = get_cost_safety_manager ()
252+ _est = 0.0001
253+ _allowed , _reason , _ = _cost_safety .check_operation (
254+ session_id = f"embed_query:{ self .expert_name } " ,
255+ operation_type = "embed_query" ,
256+ estimated_cost = _est ,
257+ require_confirmation = False ,
258+ )
259+ if not _allowed :
260+ logger .warning ("Query embedding blocked by cost-safety: %s" , _reason )
261+ return []
262+ except Exception :
263+ _cost_safety = None # type: ignore[assignment]
264+ _est = 0.0
265+
208266 # Embed query (single API call)
209267 try :
210268 response = await client .embeddings .create (model = model , input = query )
211269 query_embedding = np .array (response .data [0 ].embedding )
270+ if _cost_safety is not None :
271+ try :
272+ _cost_safety .record_cost (
273+ session_id = f"embed_query:{ self .expert_name } " ,
274+ operation_type = "embed_query" ,
275+ actual_cost = float (_est ),
276+ provider = "openai" ,
277+ model = model ,
278+ source = "experts.embedding_cache.search" ,
279+ )
280+ except Exception :
281+ pass
212282 except Exception as e :
213283 logger .error ("Error embedding query: %s" , e )
214284 return []
0 commit comments