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
2025-12-18 16:33:17 +01:00
parent 0abe12382a
commit 40d1d8cc24
5 changed files with 79 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/chat/agent_rag/document_converter/parser.py
+66
View File
@@ -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
)
src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py
+2 -55
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__)
@@ -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.
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:
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:
"""
src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py
+6 -61
View File
@@ -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,
src/backend/utils/oidc.py
+2
View File
@@ -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