From 41f694d3772df089bd5ef5da33d97c848e7cf612 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 26 Feb 2026 08:55:53 +0200 Subject: [PATCH] Better exception handling --- haiku_rag_slim/haiku/rag/client.py | 26 +++++++++----- tests/test_title_generation.py | 58 ++++++++++++++++++++++++++---- 2 files changed, 68 insertions(+), 16 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 20426bb4..1a6d5069 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -304,13 +304,9 @@ class HaikuRAG: "Return ONLY the title text, nothing else." ), ) - try: - result = await agent.run(truncated) - title = result.output.strip() - return title if title else None - except Exception: - logger.warning("LLM title generation failed", exc_info=True) - return None + result = await agent.run(truncated) + title = result.output.strip() + return title if title else None async def _resolve_title( self, @@ -335,7 +331,13 @@ class HaikuRAG: if structural: return structural - return await self._generate_title_with_llm(content) + try: + return await self._generate_title_with_llm(content) + except Exception: + logger.warning( + "LLM title generation failed during ingestion", exc_info=True + ) + return None async def generate_title(self, document: Document) -> str | None: """Generate a title for a document. @@ -1691,7 +1693,13 @@ class HaikuRAG: if doc.title is not None: continue assert doc.id is not None - title = await self.generate_title(doc) + try: + title = await self.generate_title(doc) + except Exception: + logger.warning( + "Failed to generate title for document %s", doc.id, exc_info=True + ) + continue if title is not None: doc.title = title await self.document_repository.update(doc) diff --git a/tests/test_title_generation.py b/tests/test_title_generation.py index 40474a34..d239d283 100644 --- a/tests/test_title_generation.py +++ b/tests/test_title_generation.py @@ -174,14 +174,18 @@ class TestResolveTitle: assert result == "Auto Extracted Title" @pytest.mark.asyncio - async def test_no_structural_title_no_llm_returns_none(self, tmp_path): - """Returns None when no structural title and no LLM available.""" + async def test_llm_failure_returns_none(self, tmp_path, monkeypatch): + """LLM failure during ingestion returns None instead of raising.""" doc = DoclingDocument(name="test") doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text") client = self._make_client(tmp_path) + + async def exploding_llm(self, content): + raise RuntimeError("LLM is down") + + monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm) result = await client._resolve_title(None, doc, "some content") - # LLM call will fail without allow_model_requests, so we get None assert result is None @@ -277,16 +281,15 @@ class TestGenerateTitle: @pytest.mark.asyncio async def test_no_structural_title_no_llm(self, temp_db_path): - """generate_title returns None when no structural title and LLM unavailable.""" + """generate_title raises when no structural title and LLM unavailable.""" async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document( "Just plain text without headings.", uri="test://gen-no-title", format="plain", ) - title = await client.generate_title(doc) - # LLM blocked by ALLOW_MODEL_REQUESTS=False, so falls back to None - assert title is None + with pytest.raises(RuntimeError): + await client.generate_title(doc) @pytest.mark.asyncio async def test_bypasses_auto_title_config(self, temp_db_path): @@ -357,3 +360,44 @@ class TestRebuildTitleOnly: processed_ids.append(doc_id) assert len(processed_ids) == 0 + + @pytest.mark.asyncio + async def test_continues_on_per_document_failure(self, temp_db_path, monkeypatch): + """TITLE_ONLY mode continues when generate_title fails for a document.""" + from haiku.rag.client import RebuildMode + + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc1 = await client.create_document( + "# First Heading\n\nContent.", + uri="test://first", + ) + doc2 = await client.create_document( + "# Second Heading\n\nContent.", + uri="test://second", + ) + # Clear both titles + doc1.title = None + await client.document_repository.update(doc1) + doc2.title = None + await client.document_repository.update(doc2) + + # Make generate_title fail for the first doc, succeed for the second + original = HaikuRAG.generate_title + call_count = 0 + + async def flaky_generate(self, document): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("LLM failed") + return await original(self, document) + + monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate) + + processed_ids = [] + async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY): + processed_ids.append(doc_id) + + # Only the second doc should have been processed + assert len(processed_ids) == 1