From 88bdcc2e60a5572a3313681cc077908ac456c556 Mon Sep 17 00:00:00 2001 From: charles Date: Tue, 20 Jan 2026 17:56:45 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(backend)=20handle=20deleting=20tempor?= =?UTF-8?q?ary=20collections?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I handle deleting document in temporary collection for web_search_brave_with_document_backend --- .../document_rag_backends/albert_rag_backend.py | 4 ++-- .../document_rag_backends/base_rag_backend.py | 13 ++++++++----- .../document_rag_backends/find_rag_backend.py | 17 ++++++++++++----- src/backend/chat/tools/web_search_brave.py | 13 ++++++++++--- 4 files changed, 32 insertions(+), 15 deletions(-) diff --git a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py index 3dc4000..045a8ef 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py @@ -87,7 +87,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att self.collection_id = str(response.json()["id"]) return self.collection_id - def delete_collection(self) -> None: + def delete_collection(self, **kwargs) -> None: """ Delete the current collection """ @@ -98,7 +98,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att ) response.raise_for_status() - async def adelete_collection(self) -> None: + async def adelete_collection(self, **kwargs) -> None: """ Asynchronously delete the current collection """ diff --git a/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py index db13649..ed98c1d 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py @@ -137,19 +137,20 @@ class BaseRagBackend(ABC): self.store_document(name, document_content, **kwargs) return document_content - def delete_collection(self) -> None: + @abstractmethod + def delete_collection(self, **kwargs) -> None: """ Delete the collection. This method should handle the logic to delete the collection from the backend. """ raise NotImplementedError("Must be implemented in subclass.") - async def adelete_collection(self) -> None: + async def adelete_collection(self, **kwargs) -> None: """ Delete the collection. This method should handle the logic to delete the collection from the backend. """ - return await sync_to_async(self.delete_collection)() + return await sync_to_async(self.delete_collection)(**kwargs) @abstractmethod def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults: @@ -188,7 +189,9 @@ class BaseRagBackend(ABC): @classmethod @asynccontextmanager - async def temporary_collection_async(cls, name: str, description: Optional[str] = None): + async def temporary_collection_async( + cls, name: str, description: Optional[str] = None, **kwargs + ): """Context manager for RAG backend with temporary collections.""" backend = cls() @@ -196,4 +199,4 @@ class BaseRagBackend(ABC): try: yield backend finally: - await backend.adelete_collection() + await backend.adelete_collection(**kwargs) diff --git a/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py index 68f8145..8e36716 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py @@ -42,6 +42,7 @@ class FindRagBackend(BaseRagBackend): self.api_key = settings.FIND_API_KEY self.search_endpoint = "api/v1.0/documents/search/" self.indexing_endpoint = "api/v1.0/documents/index/" + self.deleting_endpoint = "api/v1.0/documents/delete/" self.parser = AlbertParser() # Find Rag relies on Albert parser def create_collection(self, name: str, description: Optional[str] = None) -> str: @@ -51,11 +52,18 @@ class FindRagBackend(BaseRagBackend): self.collection_id = self.collection_id or str(uuid.uuid4()) return self.collection_id - def delete_collection(self) -> None: + @with_fresh_access_token + def delete_collection(self, **kwargs) -> None: """ - Deletion not available + Delete the current collection """ - logger.warning("deletion of collections is not yet supported in FindRagBackend") + response = requests.post( + urljoin(settings.FIND_API_URL, self.deleting_endpoint), + headers={"Authorization": f"Bearer {kwargs['session'].get('oidc_access_token')}"}, + json={"tags": [f"collection-{self.collection_id}"], "service": "conversations"}, + timeout=settings.FIND_API_TIMEOUT, + ) + response.raise_for_status() def store_document(self, name: str, content: str, **kwargs) -> None: """ @@ -110,12 +118,11 @@ class FindRagBackend(BaseRagBackend): RAGWebResults: The search results. """ logger.debug("search documents in Find with query '%s'", query) - response = requests.post( urljoin(settings.FIND_API_URL, self.search_endpoint), headers={"Authorization": f"Bearer {kwargs['session'].get('oidc_access_token')}"}, json={ - "q": query, + "q": query or "*", "tags": [ f"collection-{collection_id}" for collection_id in self.get_all_collection_ids() ], diff --git a/src/backend/chat/tools/web_search_brave.py b/src/backend/chat/tools/web_search_brave.py index 55b3521..cd32dd4 100644 --- a/src/backend/chat/tools/web_search_brave.py +++ b/src/backend/chat/tools/web_search_brave.py @@ -307,19 +307,26 @@ async def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> temp_collection_name = f"tmp-{uuid.uuid4()}" try: async with document_store_backend.temporary_collection_async( - temp_collection_name + temp_collection_name, session=ctx.deps.session ) as document_store: # Fetch and store all documents concurrently tasks = [ - _fetch_and_store_async(result["url"], document_store) + _fetch_and_store_async( + result["url"], + document_store, + user_sub=ctx.deps.user.sub, + session=ctx.deps.session, + ) for result in raw_search_results ] await asyncio.gather(*tasks, return_exceptions=True) # Perform RAG search rag_results = await document_store.asearch( - query, + query=query, results_count=settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER, + session=ctx.deps.session, + user_sub=ctx.deps.user.sub, ) logger.info("RAG search returned: %s", rag_results)