From 122da834d7780ed3a7aa5020e2936e8dc890e329 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 6 Feb 2026 09:53:39 +0100 Subject: [PATCH] Use resolve_document() --- haiku_rag_slim/haiku/rag/agents/rlm/runner.py | 31 +++-------------- haiku_rag_slim/haiku/rag/client.py | 31 +++++++++++++---- tests/test_client.py | 34 +++++++++++++++++++ 3 files changed, 63 insertions(+), 33 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/runner.py b/haiku_rag_slim/haiku/rag/agents/rlm/runner.py index 2fa6f3e6..b2a4c758 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/runner.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/runner.py @@ -12,7 +12,6 @@ def build_namespace( client: Any, config: Any, context: Any, loop: asyncio.AbstractEventLoop ) -> dict[str, Any]: """Build execution namespace with haiku.rag functions injected.""" - from haiku.rag.store.repositories.document import _escape_sql_string def run_async(coro: Any) -> Any: """Run async coroutine from sync context using thread-safe scheduling.""" @@ -58,37 +57,15 @@ def build_namespace( def get_document(id_or_title: str) -> str | None: async def _get() -> str | None: - doc = await client.get_document_by_id(id_or_title) - if doc: - return doc.content - safe_input = _escape_sql_string(id_or_title) - docs = await client.list_documents(filter=f"title = '{safe_input}'") - if docs and docs[0].id: - full_doc = await client.get_document_by_id(docs[0].id) - return full_doc.content if full_doc else None - docs = await client.list_documents(filter=f"uri = '{safe_input}'") - if docs and docs[0].id: - full_doc = await client.get_document_by_id(docs[0].id) - return full_doc.content if full_doc else None - return None + doc = await client.resolve_document(id_or_title) + return doc.content if doc else None return run_async(_get()) def get_docling_document(id_or_title: str) -> Any: async def _get() -> Any: - doc = await client.get_document_by_id(id_or_title) - if doc: - return doc.get_docling_document() - safe_input = _escape_sql_string(id_or_title) - docs = await client.list_documents(filter=f"title = '{safe_input}'") - if docs and docs[0].id: - full_doc = await client.get_document_by_id(docs[0].id) - return full_doc.get_docling_document() if full_doc else None - docs = await client.list_documents(filter=f"uri = '{safe_input}'") - if docs and docs[0].id: - full_doc = await client.get_document_by_id(docs[0].id) - return full_doc.get_docling_document() if full_doc else None - return None + doc = await client.resolve_document(id_or_title) + return doc.get_docling_document() if doc else None return run_async(_get()) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 37f31bfb..67a6a349 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -739,6 +739,30 @@ class HaikuRAG: """ return await self.document_repository.get_by_uri(uri) + async def resolve_document(self, id_or_title: str) -> Document | None: + """Resolve a document by ID, title, or URI (in that order). + + Args: + id_or_title: Document ID, title, or URI to look up. + + Returns: + The Document instance if found, None otherwise. + """ + doc = await self.get_document_by_id(id_or_title) + if doc: + return doc + + safe_input = _escape_sql_string(id_or_title) + docs = await self.list_documents(filter=f"title = '{safe_input}'") + if docs and docs[0].id: + return await self.get_document_by_id(docs[0].id) + + docs = await self.list_documents(filter=f"uri = '{safe_input}'") + if docs and docs[0].id: + return await self.get_document_by_id(docs[0].id) + + return None + async def update_document( self, document_id: str, @@ -1328,12 +1352,7 @@ class HaikuRAG: if documents: loaded_docs = [] for doc_ref in documents: - doc = await self.get_document_by_id(doc_ref) - if not doc: - safe_ref = _escape_sql_string(doc_ref) - docs = await self.list_documents(filter=f"title = '{safe_ref}'") - if docs and docs[0].id: - doc = await self.get_document_by_id(docs[0].id) + doc = await self.resolve_document(doc_ref) if doc: loaded_docs.append(doc) context.documents = loaded_docs if loaded_docs else None diff --git a/tests/test_client.py b/tests/test_client.py index 33644f79..3df0d286 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -85,6 +85,40 @@ async def test_client_document_crud(qa_corpus: Dataset, temp_db_path): assert deleted_again is False +async def test_client_resolve_document(temp_db_path): + """Test resolve_document finds documents by ID, title, or URI.""" + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.create_document( + content="Test content", + uri="test://resolve-test", + title="Resolve Test Doc", + ) + + # Resolve by ID + by_id = await client.resolve_document(doc.id) + assert by_id is not None + assert by_id.id == doc.id + + # Resolve by title + by_title = await client.resolve_document("Resolve Test Doc") + assert by_title is not None + assert by_title.id == doc.id + + # Resolve by URI + by_uri = await client.resolve_document("test://resolve-test") + assert by_uri is not None + assert by_uri.id == doc.id + + # Not found returns None + not_found = await client.resolve_document("nonexistent") + assert not_found is None + + # SQL injection is escaped + injection = "x' OR title LIKE '%" + injected = await client.resolve_document(injection) + assert injected is None + + @pytest.mark.vcr() async def test_client_update_document(qa_corpus: Dataset, temp_db_path): """Test updating document with individual parameters."""