electron-rare/conversations
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

♻️(backend) refactor document parsers

Browse Source 
I refactor document parsing by introducing
AlbertParser and BaseParser
This commit is contained in:
charles
2026-02-02 22:27:26 +01:00
committed by Quentin BEY
parent 3106d5f25f
commit 1c573ad3a7
4 changed files with 82 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/chat/agent_rag/document_converter/parser.py
+65
View File
@@ -0,0 +1,65 @@
"""Document parsers for RAG backends."""
import logging
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: bytes) -> 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 (bytes): The content of the document as a bytes 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: bytes) -> 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: bytes) -> 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
)
src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py
+3 -58
View File
@@ -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__)
@@ -26,9 +26,6 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
It provides methods to:
- Create a collection for the search operation.
- Parse documents and convert them to Markdown format:
+ Handle PDF parsing using the Albert API.
+ Use the DocumentConverter (markitdown) for other formats.
- Store parsed documents in the Albert collection.
- Perform a search operation using the Albert API.
"""
@@ -46,10 +43,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 +110,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.
@@ -266,6 +210,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
Args:
query (str): The search query.
results_count (int): The number of results to return.
**kwargs: Additional arguments.
Returns:
RAGWebResults: The search results.
src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py
+3 -1
View File
@@ -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(ABC):
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]:
"""
@@ -89,7 +91,7 @@ class BaseRagBackend(ABC):
Returns:
str: The document content in Markdown format.
"""
raise NotImplementedError("Must be implemented in subclass.")
return self.parser.parse_document(name, content_type, content)
@abstractmethod
def store_document(self, name: str, content: str, **kwargs) -> None:
src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py
+11 -70
View File
@@ -2,34 +2,32 @@
import logging
import uuid
from io import BytesIO
from typing import List, Optional
from urllib.parse import urljoin
from uuid import uuid4
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
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_converter.parser import AlbertParser
from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend
from utils.oidc import with_fresh_access_token
logger = logging.getLogger(__name__)
class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attributes
SUPPORTED_LANGUAGE_CODES = ["en", "fr", "de", "nl"]
class FindRagBackend(BaseRagBackend):
"""
This class is a placeholder for the Find API implementation.
It is designed to be used with the RAG (Retrieval-Augmented Generation) document search system.
It provides methods to:
- Parse documents and convert them to Markdown format:
+ Handle PDF parsing using the Albert API.
+ Use the DocumentConverter (markitdown) for other formats.
- Store parsed documents in the Find index.
- Perform a search operation using the Find API.
"""
@@ -41,15 +39,12 @@ 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() # Find Rag relies on Albert parser
if not self.api_key:
raise ImproperlyConfigured("FIND_API_KEY must be set in Django settings.")
def create_collection(self, name: str, description: Optional[str] = None) -> uuid.UUID:
def create_collection(self, name: str, description: Optional[str] = None) -> str:
"""
init collection_id
"""
@@ -62,62 +57,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
@@ -174,10 +113,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,