From 40d1d8cc2487537ae3ff3e766be3d5a442a566e5 Mon Sep 17 00:00:00 2001 From: charles Date: Thu, 18 Dec 2025 16:33:17 +0100 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F(backend)=20refactor=20docume?= =?UTF-8?q?nt=20parsers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I refactor document parsing by introducing AlbertParser and BaseParser --- .../agent_rag/document_converter/parser.py | 66 ++++++++++++++++++ .../albert_rag_backend.py | 57 +--------------- .../document_rag_backends/base_rag_backend.py | 4 +- .../document_rag_backends/find_rag_backend.py | 67 ++----------------- src/backend/utils/oidc.py | 2 + 5 files changed, 79 insertions(+), 117 deletions(-) create mode 100644 src/backend/chat/agent_rag/document_converter/parser.py diff --git a/src/backend/chat/agent_rag/document_converter/parser.py b/src/backend/chat/agent_rag/document_converter/parser.py new file mode 100644 index 0000000..d182046 --- /dev/null +++ b/src/backend/chat/agent_rag/document_converter/parser.py @@ -0,0 +1,66 @@ +"""Document parsers for RAG backends.""" + +import logging +from io import BytesIO +from urllib.parse import urljoin + +from django.conf import settings + +import requests + +from chat.agent_rag.document_converter.markitdown import DocumentConverter + +logger = logging.getLogger(__name__) + + +class BaseParser: + """Base class for document parsers.""" + + def parse_document(self, name: str, content_type: str, content: BytesIO) -> str: + """ + Parse the document and prepare it for the search operation. + This method should handle the logic to convert the document + into a format suitable for storage. + + Args: + name (str): The name of the document. + content_type (str): The MIME type of the document (e.g., "application/pdf"). + content (BytesIO): The content of the document as a BytesIO stream. + + Returns: + str: The document content in Markdown format. + """ + raise NotImplementedError("Must be implemented in subclass.") + + +class AlbertParser(BaseParser): + """Document parser using Albert API for PDFs and DocumentConverter for other formats.""" + + endpoint = urljoin(settings.ALBERT_API_URL, "/v1/parse-beta") + + def parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str: + """Parse PDF document using Albert API.""" + response = requests.post( + self.endpoint, + headers={ + "Authorization": f"Bearer {settings.ALBERT_API_KEY}", + }, + files={ + "file": (name, content, content_type), + "output_format": (None, "markdown"), + }, + timeout=settings.ALBERT_API_PARSE_TIMEOUT, + ) + response.raise_for_status() + + return "\n\n".join( + document_page["content"] for document_page in response.json().get("data", []) + ) + + def parse_document(self, name: str, content_type: str, content: BytesIO) -> str: + """Parse document based on content type.""" + if content_type == "application/pdf": + return self.parse_pdf_document(name=name, content_type=content_type, content=content) + return DocumentConverter().convert_raw( + name=name, content_type=content_type, content=content + ) 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 9753f4a..562201c 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 @@ -13,7 +13,7 @@ import requests from chat.agent_rag.albert_api_constants import Searches from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage -from chat.agent_rag.document_converter.markitdown import DocumentConverter +from chat.agent_rag.document_converter.parser import AlbertParser from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend logger = logging.getLogger(__name__) @@ -46,10 +46,9 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att } self._collections_endpoint = urljoin(self._base_url, "/v1/collections") self._documents_endpoint = urljoin(self._base_url, "/v1/documents") - self._pdf_parser_endpoint = urljoin(self._base_url, "/v1/parse-beta") self._search_endpoint = urljoin(self._base_url, "/v1/search") - self._default_collection_description = "Temporary collection for RAG document search" + self.parser = AlbertParser() def create_collection(self, name: str, description: Optional[str] = None) -> str: """ @@ -114,58 +113,6 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att ) response.raise_for_status() - def parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str: - """ - Parse the PDF document content and return the text content. - This method should handle the logic to convert the PDF into - a format suitable for the Albert API. - """ - response = requests.post( - self._pdf_parser_endpoint, - headers=self._headers, - files={ - "file": ( - name, - content, - content_type, - ), # Use the name as the filename in the request - "output_format": (None, "markdown"), # Specify the output format as Markdown, - }, - timeout=settings.ALBERT_API_PARSE_TIMEOUT, - ) - response.raise_for_status() - - return "\n\n".join( - document_page["content"] for document_page in response.json().get("data", []) - ) - - def parse_document(self, name: str, content_type: str, content: BytesIO): - """ - Parse the document and prepare it for the search operation. - This method should handle the logic to convert the document - into a format suitable for the Albert API. - - Args: - name (str): The name of the document. - content_type (str): The MIME type of the document (e.g., "application/pdf"). - content (BytesIO): The content of the document as a BytesIO stream. - - Returns: - str: The document content in Markdown format. - """ - # Implement the parsing logic here - if content_type == "application/pdf": - # Handle PDF parsing - markdown_content = self.parse_pdf_document( - name=name, content_type=content_type, content=content - ) - else: - markdown_content = DocumentConverter().convert_raw( - name=name, content_type=content_type, content=content - ) - - return markdown_content - def store_document(self, name: str, content: str, **kwargs) -> None: """ Store the document content in the Albert 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 4eb185f..d658869 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 @@ -8,6 +8,7 @@ from typing import List, Optional from asgiref.sync import sync_to_async from chat.agent_rag.constants import RAGWebResults +from chat.agent_rag.document_converter.parser import BaseParser logger = logging.getLogger(__name__) @@ -38,6 +39,7 @@ class BaseRagBackend: self.collection_id = collection_id self.read_only_collection_id = read_only_collection_id or [] self._default_collection_description = "Temporary collection for RAG document search" + self.parser: BaseParser = BaseParser() def get_all_collection_ids(self) -> List[str]: """ @@ -88,7 +90,7 @@ class BaseRagBackend: Returns: str: The document content in Markdown format. """ - raise NotImplementedError("Must be implemented in subclass.") + return self.parser.parse_document(name, content_type, content) def store_document(self, name: str, content: str, **kwargs) -> None: """ 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 4197212..2780dc4 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 @@ -2,7 +2,6 @@ import logging import uuid -from io import BytesIO from typing import List, Optional from urllib.parse import urljoin from uuid import uuid4 @@ -14,7 +13,7 @@ from django.utils import timezone import requests from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage -from chat.agent_rag.document_converter.markitdown import DocumentConverter +from chat.agent_rag.document_rag_backends.albert_rag_backend import AlbertParser from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend from utils.oidc import with_fresh_access_token @@ -41,10 +40,10 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri ): # Initialize any necessary parameters or configurations here super().__init__(collection_id, read_only_collection_id) - self._pdf_parser_endpoint = urljoin(settings.ALBERT_API_URL, "/v1/parse-beta") self.api_key = settings.FIND_API_KEY self.search_endpoint = "api/v1.0/documents/search/" self.indexing_endpoint = "api/v1.0/documents/index/" + self.parser = AlbertParser() if not self.api_key: raise ImproperlyConfigured("FIND_API_KEY must be set in Django settings.") @@ -62,62 +61,6 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri """ logger.warning("deletion of collections is not yet supported in FindRagBackend") - # TODO: factor with albert api - def parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str: - """ - Parse the PDF document content and return the text content. - This method should handle the logic to convert the PDF into - a format suitable for the Albert API. - """ - response = requests.post( - self._pdf_parser_endpoint, - headers={ - "Authorization": f"Bearer {settings.ALBERT_API_KEY}", - }, - files={ - "file": ( - name, - content, - content_type, - ), # Use the name as the filename in the request - "output_format": (None, "markdown"), # Specify the output format as Markdown, - }, - timeout=settings.ALBERT_API_PARSE_TIMEOUT, - ) - response.raise_for_status() - - return "\n\n".join( - document_page["content"] for document_page in response.json().get("data", []) - ) - - # TODO: factor with albert api - def parse_document(self, name: str, content_type: str, content: BytesIO): - """ - Parse the document and prepare it for the search operation. - This method should handle the logic to convert the document - into a format suitable for the Find API. - - Args: - name (str): The name of the document. - content_type (str): The MIME type of the document (e.g., "application/pdf"). - content (BytesIO): The content of the document as a BytesIO stream. - - Returns: - str: The document content in Markdown format. - """ - # Implement the parsing logic here - if content_type == "application/pdf": - # Handle PDF parsing - markdown_content = self.parse_pdf_document( - name=name, content_type=content_type, content=content - ) - else: - markdown_content = DocumentConverter().convert_raw( - name=name, content_type=content_type, content=content - ) - - return markdown_content - def store_document(self, name: str, content: str, **kwargs) -> None: """ index document in Find @@ -169,10 +112,12 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri response = requests.post( urljoin(settings.FIND_API_URL, self.search_endpoint), - headers={"Authorization": f"Bearer {kwargs["session"].get("oidc_access_token")}"}, + headers={"Authorization": f"Bearer {kwargs['session'].get('oidc_access_token')}"}, json={ "q": query, - "tags": [f"collection-{collection_id}" for collection_id in self.get_all_collection_ids()], + "tags": [ + f"collection-{collection_id}" for collection_id in self.get_all_collection_ids() + ], "k": results_count, }, timeout=settings.FIND_API_TIMEOUT, diff --git a/src/backend/utils/oidc.py b/src/backend/utils/oidc.py index 231e2a6..6bba396 100644 --- a/src/backend/utils/oidc.py +++ b/src/backend/utils/oidc.py @@ -38,6 +38,7 @@ def with_fresh_access_token(func): Decorator to handle OIDC token refresh and extraction. Expects 'session' in kwargs and update it with the fresh token. """ + @wraps(func) def wrapper(*args, **kwargs): session = kwargs.get("session") @@ -45,4 +46,5 @@ def with_fresh_access_token(func): raise AuthenticationFailed({"error": "Session is required but not provided"}) kwargs["session"] = refresh_access_token(session) return func(*args, **kwargs) + return wrapper