Skip to content

Commit 37740da

Browse files
thezfzthezfz
committed
faiss_error
1 parent ad4474b commit 37740da

File tree

2 files changed

+90
-106
lines changed

2 files changed

+90
-106
lines changed

‎aigraphx/core/db.py‎

Lines changed: 58 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -115,71 +115,89 @@ async def lifespan(
115115

116116
# Initialize Faiss Repository for Papers (Simplified Error Handling)
117117
logger.info("Initializing Faiss Repository for Papers...")
118+
faiss_repo_papers_instance = None # Temporary variable
118119
try:
119-
app.state.faiss_repo_papers = FaissRepository(
120+
# Only try to instantiate the repository here
121+
faiss_repo_papers_instance = FaissRepository(
120122
index_path=settings.faiss_index_path,
121123
id_map_path=settings.faiss_mapping_path,
122124
id_type="int",
123125
)
124-
# Check readiness immediately after instantiation
125-
if not app.state.faiss_repo_papers.is_ready():
126-
logger.error("Papers Faiss Repository is not ready after initialization.")
127-
# Log the state from is_ready for more info
128-
index_exists = app.state.faiss_repo_papers.index is not None
129-
ntotal = (
130-
app.state.faiss_repo_papers.index.ntotal
131-
if app.state.faiss_repo_papers.index
132-
else -1
133-
)
134-
map_exists_and_not_empty = bool(app.state.faiss_repo_papers.id_map)
126+
app.state.faiss_repo_papers = faiss_repo_papers_instance # Assign if instantiation succeeded
127+
logger.info("FaissRepository for Papers instantiated.")
128+
129+
except Exception as e:
130+
# Catch only instantiation errors
131+
logger.exception(f"Failed to initialize Papers Faiss Repository (Instantiation Error): {e}")
132+
app.state.faiss_repo_papers = None # Ensure state is None on instantiation exception
133+
logger.error("CRITICAL: Papers Faiss Repository instantiation failed. Raising RuntimeError.")
134+
raise RuntimeError("Papers Faiss Repository initialization failed") from e
135+
136+
# Now, check readiness *outside* the instantiation try...except block
137+
# Only proceed if instantiation was successful (faiss_repo_papers_instance is not None)
138+
if faiss_repo_papers_instance is not None:
139+
logger.info("Checking readiness of Papers Faiss Repository...")
140+
if not faiss_repo_papers_instance.is_ready():
141+
index_exists = os.path.exists(settings.faiss_index_path)
142+
ntotal = faiss_repo_papers_instance.index.ntotal if faiss_repo_papers_instance.index else None
143+
map_exists_and_not_empty = bool(faiss_repo_papers_instance.id_map)
135144
logger.error(
136145
f"[Lifespan Check Failed] Papers Faiss State: index_exists={index_exists}, ntotal={ntotal}, map_not_empty={map_exists_and_not_empty}"
137146
)
138-
# Set state to None if not ready
147+
# Raise the specific error for not being ready
148+
logger.error("CRITICAL: Papers Faiss Repository is not ready. Raising RuntimeError.")
149+
# Set state to None *before* raising, as the repo is unusable
139150
app.state.faiss_repo_papers = None
140-
# Do not raise RuntimeError here to allow tests to check state
141-
logger.warning(
142-
"Papers Faiss Repository set to None due to initialization failure."
143-
)
151+
raise RuntimeError("Papers Faiss Repository is not ready after initialization.")
144152
else:
145153
logger.info("Papers Faiss Repository initialized and ready.")
146-
except Exception as e:
147-
logger.exception(f"Failed to initialize Papers Faiss Repository: {e}")
148-
app.state.faiss_repo_papers = None # Ensure state is None on exception
154+
else:
155+
# This case is technically handled by the raise in the except block,
156+
# but adding an info log might be useful for clarity.
157+
logger.info("Skipping readiness check for Papers Faiss Repository due to instantiation failure.")
158+
149159

150160
# Initialize Faiss Repository for Models (Simplified Error Handling)
151161
logger.info("Initializing Faiss Repository for Models...")
162+
faiss_repo_models_instance = None # Temporary variable
152163
try:
153-
app.state.faiss_repo_models = FaissRepository(
164+
# Only try to instantiate the repository here
165+
faiss_repo_models_instance = FaissRepository(
154166
index_path=settings.models_faiss_index_path,
155167
id_map_path=settings.models_faiss_mapping_path,
156168
id_type="str",
157169
)
158-
# Check readiness immediately after instantiation
159-
if not app.state.faiss_repo_models.is_ready():
160-
logger.error("Models Faiss Repository is not ready after initialization.")
161-
# Log the state from is_ready for more info
162-
index_exists_m = app.state.faiss_repo_models.index is not None
163-
ntotal_m = (
164-
app.state.faiss_repo_models.index.ntotal
165-
if app.state.faiss_repo_models.index
166-
else -1
167-
)
168-
map_exists_and_not_empty_m = bool(app.state.faiss_repo_models.id_map)
170+
app.state.faiss_repo_models = faiss_repo_models_instance # Assign if instantiation succeeded
171+
logger.info("FaissRepository for Models instantiated.")
172+
173+
except Exception as e:
174+
# Catch only instantiation errors
175+
logger.exception(f"Failed to initialize Models Faiss Repository (Instantiation Error): {e}")
176+
app.state.faiss_repo_models = None # Ensure state is None on instantiation exception
177+
logger.error("CRITICAL: Models Faiss Repository instantiation failed. Raising RuntimeError.")
178+
raise RuntimeError("Models Faiss Repository initialization failed") from e
179+
180+
# Now, check readiness *outside* the instantiation try...except block
181+
# Only proceed if instantiation was successful
182+
if faiss_repo_models_instance is not None:
183+
logger.info("Checking readiness of Models Faiss Repository...")
184+
if not faiss_repo_models_instance.is_ready():
185+
index_exists_m = os.path.exists(settings.models_faiss_index_path)
186+
ntotal_m = faiss_repo_models_instance.index.ntotal if faiss_repo_models_instance.index else None
187+
map_exists_and_not_empty_m = bool(faiss_repo_models_instance.id_map)
169188
logger.error(
170189
f"[Lifespan Check Failed] Models Faiss State: index_exists={index_exists_m}, ntotal={ntotal_m}, map_not_empty={map_exists_and_not_empty_m}"
171190
)
172-
# Set state to None if not ready
191+
# Raise the specific error for not being ready
192+
logger.error("CRITICAL: Models Faiss Repository is not ready. Raising RuntimeError.")
193+
# Set state to None *before* raising
173194
app.state.faiss_repo_models = None
174-
# Do not raise RuntimeError here to allow tests to check state
175-
logger.warning(
176-
"Models Faiss Repository set to None due to initialization failure."
177-
)
195+
raise RuntimeError("Models Faiss Repository is not ready after initialization.")
178196
else:
179197
logger.info("Models Faiss Repository initialized and ready.")
180-
except Exception as e:
181-
logger.exception(f"Failed to initialize Models Faiss Repository: {e}")
182-
app.state.faiss_repo_models = None # Ensure state is None on exception
198+
else:
199+
logger.info("Skipping readiness check for Models Faiss Repository due to instantiation failure.")
200+
183201

184202
logger.info("Resource initialization process completed.")
185203
print("--- DEBUG: Lifespan YIELDING (Resources should be initialized) ---")

‎tests/core/test_db.py‎

Lines changed: 32 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,12 @@ def mock_faiss_repository() -> Generator[
5656
# Create two distinct mock instances for papers and models
5757
mock_instance_papers = MagicMock(spec=FaissRepository)
5858
mock_instance_papers.index = MagicMock()
59+
mock_instance_papers.id_map = {} # Add id_map attribute
5960
mock_instance_papers.is_ready.return_value = True # Default to ready
6061

6162
mock_instance_models = MagicMock(spec=FaissRepository)
6263
mock_instance_models.index = MagicMock()
64+
mock_instance_models.id_map = {} # Add id_map attribute
6365
mock_instance_models.is_ready.return_value = True # Default to ready
6466

6567
# Configure the class mock to return the correct instance based on call args
@@ -252,44 +254,25 @@ async def test_lifespan_faiss_not_ready(
252254
mock_faiss_repository
253255
)
254256

255-
# Simulate Faiss Papers repo being not ready after init
257+
# Simulate Papers Faiss repo being not ready
256258
mock_repo_instance_papers.is_ready.return_value = False
257-
mock_repo_instance_models.is_ready.return_value = True # Models repo is fine
259+
mock_repo_instance_models.is_ready.return_value = True # Models is ready
258260

259261
ctx = lifespan(mock_app, mock_settings)
260-
# Lifespan should NOT raise an exception here, just log warning and set state to None
261-
await ctx.__aenter__()
262-
263-
# Assertions after startup
264-
mock_pool_class.assert_called_once()
265-
mock_driver_func.assert_called_once()
266-
mock_repo_class.assert_any_call(
267-
index_path=mock_settings.faiss_index_path,
268-
id_map_path=mock_settings.faiss_mapping_path,
269-
id_type="int",
270-
)
271-
mock_repo_class.assert_any_call(
272-
index_path=mock_settings.models_faiss_index_path,
273-
id_map_path=mock_settings.models_faiss_mapping_path,
274-
id_type="str",
275-
)
276-
assert mock_repo_class.call_count == 2
277-
mock_repo_instance_papers.is_ready.assert_called_once() # is_ready was checked
278-
mock_repo_instance_models.is_ready.assert_called_once() # is_ready was checked
262+
with pytest.raises(RuntimeError, match="Papers Faiss Repository is not ready after initialization."):
263+
await ctx.__aenter__()
279264

280-
assert mock_app.state.pg_pool == mock_pool_instance
281-
assert mock_app.state.neo4j_driver == mock_driver_instance
282-
assert (
283-
mock_app.state.faiss_repo_papers is None
284-
) # State should be None as it wasn't ready
265+
# Check state after expected failure
266+
assert mock_app.state.pg_pool == mock_pool_instance # Should still be set
267+
assert mock_app.state.neo4j_driver == mock_driver_instance # Should still be set
268+
assert mock_app.state.faiss_repo_papers is None # Should be None or not set
285269
assert (
286-
mock_app.state.faiss_repo_models == mock_repo_instance_models
287-
) # Models repo is fine
270+
mock_app.state.faiss_repo_models is None
271+
) # Should also be None as init likely stopped
288272

289-
# Check shutdown calls
290-
await ctx.__aexit__(None, None, None)
291-
mock_pool_instance.close.assert_awaited_once()
292-
mock_driver_instance.close.assert_awaited_once()
273+
# Check close calls are not made if startup failed mid-way
274+
mock_pool_instance.close.assert_not_awaited()
275+
mock_driver_instance.close.assert_not_awaited()
293276

294277

295278
# Test initialization failure (e.g., Faiss Papers Repo __init__ fails)
@@ -320,27 +303,18 @@ def side_effect(*args: Any, **kwargs: Any) -> Any:
320303
mock_repo_class.side_effect = side_effect
321304

322305
ctx = lifespan(mock_app, mock_settings)
323-
# Lifespan should NOT raise error, just log and set state to None
324-
await ctx.__aenter__()
325-
326-
# Assertions after startup
327-
mock_pool_class.assert_called_once()
328-
mock_driver_func.assert_called_once()
329-
assert mock_repo_class.call_count == 2 # Both init attempts were made
330-
mock_repo_instance_papers.is_ready.assert_not_called() # is_ready not called if init fails
331-
mock_repo_instance_models.is_ready.assert_called_once() # is_ready called for models
306+
with pytest.raises(RuntimeError, match="Papers Faiss Repository initialization failed"):
307+
await ctx.__aenter__()
332308

309+
# Check state after expected failure
333310
assert mock_app.state.pg_pool == mock_pool_instance
334311
assert mock_app.state.neo4j_driver == mock_driver_instance
335-
assert mock_app.state.faiss_repo_papers is None # State is None due to init failure
336-
assert (
337-
mock_app.state.faiss_repo_models == mock_repo_instance_models
338-
) # Models repo ok
312+
assert mock_app.state.faiss_repo_papers is None
313+
assert mock_app.state.faiss_repo_models is None
339314

340-
# Check shutdown calls
341-
await ctx.__aexit__(None, None, None)
342-
mock_pool_instance.close.assert_awaited_once()
343-
mock_driver_instance.close.assert_awaited_once()
315+
# Check close calls
316+
mock_pool_instance.close.assert_not_awaited()
317+
mock_driver_instance.close.assert_not_awaited()
344318

345319

346320
# Test initialization failure (e.g., Faiss Models Repo __init__ fails)
@@ -371,27 +345,19 @@ def side_effect(*args: Any, **kwargs: Any) -> Any:
371345
mock_repo_class.side_effect = side_effect
372346

373347
ctx = lifespan(mock_app, mock_settings)
374-
# Lifespan should NOT raise error
375-
await ctx.__aenter__()
376-
377-
# Assertions after startup
378-
mock_pool_class.assert_called_once()
379-
mock_driver_func.assert_called_once()
380-
assert mock_repo_class.call_count == 2
381-
mock_repo_instance_papers.is_ready.assert_called_once() # Papers is_ready called
382-
mock_repo_instance_models.is_ready.assert_not_called() # Models is_ready not called
348+
with pytest.raises(RuntimeError, match="Models Faiss Repository initialization failed"):
349+
await ctx.__aenter__()
383350

351+
# Check state after expected failure
384352
assert mock_app.state.pg_pool == mock_pool_instance
385353
assert mock_app.state.neo4j_driver == mock_driver_instance
386-
assert (
387-
mock_app.state.faiss_repo_papers == mock_repo_instance_papers
388-
) # Papers repo ok
389-
assert mock_app.state.faiss_repo_models is None # State is None due to init failure
354+
# Paper repo might be initialized before model repo fails
355+
assert mock_app.state.faiss_repo_papers == mock_repo_instance_papers
356+
assert mock_app.state.faiss_repo_models is None
390357

391-
# Check shutdown calls
392-
await ctx.__aexit__(None, None, None)
393-
mock_pool_instance.close.assert_awaited_once()
394-
mock_driver_instance.close.assert_awaited_once()
358+
# Check close calls
359+
mock_pool_instance.close.assert_not_awaited()
360+
mock_driver_instance.close.assert_not_awaited()
395361

396362

397363
# Test Neo4j not configured scenario

0 commit comments

Comments
 (0)