diff --git a/apps/reports/ai_features/extraction.py b/apps/reports/ai_features/extraction.py index 44efa82..8a9ec3f 100644 --- a/apps/reports/ai_features/extraction.py +++ b/apps/reports/ai_features/extraction.py @@ -15,9 +15,8 @@ @dataclass -class PdfExtraction: +class BaseExtraction: report: Report - data: bytes llm_handler: OllamaHandler = field(init=False) llm_chat_model: ChatOllama | ChatOpenAI = field(init=False) @@ -31,6 +30,34 @@ def __post_init__(self): except Exception as e: raise e + def handle_meta_info(self): + """Extract meta information of the report.""" + title = self.report.title + description = self.report.description + if title and title.strip(): + DocumentExtraction.objects.create( + report=self.report, + status=DocumentExtractionStatus.SUCCESS, + text=title, + page_number=None, + chunk_type=DocumentExtraction.ExtractionType.TITLE, + embedding=self.llm_embedding_model.embed_query(title), + ) + if description and description.strip(): + DocumentExtraction.objects.create( + report=self.report, + status=DocumentExtractionStatus.SUCCESS, + text=description, + page_number=None, + chunk_type=DocumentExtraction.ExtractionType.DESCRIPTION, + embedding=self.llm_embedding_model.embed_query(description), + ) + + +@dataclass +class PdfExtraction(BaseExtraction): + data: bytes + def img_to_base64(self, data: fitz.Pixmap) -> str: img_bytes = data.tobytes("png") return base64.b64encode(img_bytes).decode("utf-8") @@ -39,6 +66,9 @@ def pdf_to_images(self, zoom: float = 2.0): page_summaries = [] doc = fitz.open(stream=self.data, filetype="pdf") + # Get the title and description + self.handle_meta_info() + for page_idx in range(len(doc)): page = doc[page_idx] @@ -108,3 +138,8 @@ def pdf_to_images(self, zoom: float = 2.0): ) else: logger.warning("The key doc_summary is missing in the output") + + +@dataclass +class HeaderExtraction(BaseExtraction): + """Get the basic meta information extraction.""" diff --git a/apps/reports/ai_features/search_docs.py b/apps/reports/ai_features/search_docs.py index 415a238..f492e76 100644 --- a/apps/reports/ai_features/search_docs.py +++ b/apps/reports/ai_features/search_docs.py @@ -35,6 +35,8 @@ def __post_init__(self): DocumentExtraction.ExtractionType.KEYWORDS: 0.9, DocumentExtraction.ExtractionType.TABLE: 0.3, DocumentExtraction.ExtractionType.CHART: 0.3, + DocumentExtraction.ExtractionType.TITLE: 0.2, + DocumentExtraction.ExtractionType.DESCRIPTION: 0.2, } def generate_query_embedding(self) -> list[float]: diff --git a/apps/reports/extraction.py b/apps/reports/extraction.py index 75f8230..bf49698 100644 --- a/apps/reports/extraction.py +++ b/apps/reports/extraction.py @@ -2,7 +2,7 @@ import logging -from apps.reports.models import Report +from apps.reports.models import Report, ReportContentType from apps.reports.tasks.task import handle_documents logger = logging.getLogger(__name__) @@ -15,13 +15,12 @@ async def trigger_document_extraction(report: Report) -> None: The AI tool is expected to update extracted_contents, search_text, summary, and status directly in the database once processing completes. """ - if not report.file.name: + if report.content_type == ReportContentType.FILE and not report.file.name: logger.warning("Report file is missing") return logger.info( - "Triggering document extraction for report_id=%s file_path=%s", + "Triggering document extraction for report_id=%s", str(report.pk), - report.file.name, ) - handle_documents.apply_async(args=[report.pk]) # pyright: ignore[reportFunctionMemberAccess] + handle_documents.apply_async(args=(report.pk, report.content_type)) # pyright: ignore[reportFunctionMemberAccess] diff --git a/apps/reports/migrations/0011_alter_documentextraction_chunk_type.py b/apps/reports/migrations/0011_alter_documentextraction_chunk_type.py new file mode 100644 index 0000000..fb1d781 --- /dev/null +++ b/apps/reports/migrations/0011_alter_documentextraction_chunk_type.py @@ -0,0 +1,19 @@ +# Generated by Django 6.0.3 on 2026-06-30 09:57 + +import django_choices_field.fields +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('reports', '0010_alter_report_thematic_area'), + ] + + operations = [ + migrations.AlterField( + model_name='documentextraction', + name='chunk_type', + field=django_choices_field.fields.IntegerChoicesField(choices=[(1, 'Extracted Content'), (2, 'Document Summary'), (3, 'Page Summary'), (4, 'Keywords'), (5, 'Table'), (6, 'Chart'), (7, 'Title'), (8, 'Description')], default=2), + ), + ] diff --git a/apps/reports/models.py b/apps/reports/models.py index 377ef08..5f333ac 100644 --- a/apps/reports/models.py +++ b/apps/reports/models.py @@ -166,6 +166,8 @@ class ExtractionType(models.IntegerChoices): KEYWORDS = 4, "Keywords" TABLE = 5, "Table" CHART = 6, "Chart" + TITLE = 7, "Title" + DESCRIPTION = 8, "Description" report = models.ForeignKey( Report, diff --git a/apps/reports/tasks/task.py b/apps/reports/tasks/task.py index adef113..86de0de 100644 --- a/apps/reports/tasks/task.py +++ b/apps/reports/tasks/task.py @@ -3,25 +3,33 @@ from celery import shared_task from django.core.files.storage import default_storage -from apps.reports.ai_features.extraction import PdfExtraction -from apps.reports.models import Report +from apps.reports.ai_features.extraction import HeaderExtraction, PdfExtraction +from apps.reports.models import Report, ReportContentType logger = logging.getLogger(__name__) @shared_task -def handle_documents(report_id: int) -> None: +def handle_documents(report_id: int, report_type: ReportContentType) -> None: report = Report.objects.filter(pk=report_id).first() - if report is None or report.file.name is None: - logger.warning("Report with id %s is set to None", report_id) + if report is None: + logger.warning("Report with id %s not found", report_id) return - with default_storage.open(report.file.name, "rb") as f: - pdf_bytes = f.read() + if report_type == ReportContentType.FILE and report.file.name is None: + logger.warning("Report with id %s has no file", report_id) + return - ext = PdfExtraction( - report=report, - data=pdf_bytes, - ) - ext.pdf_to_images() + if report.file and (file_name := report.file.name) and report_type == ReportContentType.FILE: + with default_storage.open(file_name, "rb") as f: + pdf_bytes = f.read() + + ext = PdfExtraction( + report=report, + data=pdf_bytes, + ) + ext.pdf_to_images() + elif report_type == ReportContentType.IFRAME: + ext = HeaderExtraction(report=report) + ext.handle_meta_info() diff --git a/apps/reports/tests/mutations_test.py b/apps/reports/tests/mutations_test.py index 8e263f5..28f3329 100644 --- a/apps/reports/tests/mutations_test.py +++ b/apps/reports/tests/mutations_test.py @@ -111,23 +111,26 @@ def setUpClass(cls): def test_create_iframe_report(self): self.force_login(self.user) - content = self.query_check( - self.Mutation.CREATE_REPORT, - variables={ - "data": { - "title": "Test Report", - "contentType": Report.ContentType.IFRAME.name, - "iframeUrl": "https://app.powerbi.com/embed/123", - "thematicArea": str(self.thematic_area.pk), + with patch("apps.reports.graphql.mutations.trigger_document_extraction") as mock_trigger: + mock_trigger.return_value = None + content = self.query_check( + self.Mutation.CREATE_REPORT, + variables={ + "data": { + "title": "Test Report", + "contentType": Report.ContentType.IFRAME.name, + "iframeUrl": "https://app.powerbi.com/embed/123", + "thematicArea": str(self.thematic_area.pk), + }, }, - }, - ) + ) resp = content["data"]["createReport"] assert resp["ok"] is True assert resp["errors"] is None assert resp["result"]["title"] == "Test Report" assert resp["result"]["contentType"] == Report.ContentType.IFRAME.name assert resp["result"]["visibility"] == Report.Visibility.PUBLIC.name + mock_trigger.assert_called_once() def test_create_report_requires_auth(self): self.logout() diff --git a/schema.graphql b/schema.graphql index 84fdc45..66b6861 100644 --- a/schema.graphql +++ b/schema.graphql @@ -462,6 +462,8 @@ enum ExtractionType { KEYWORDS TABLE CHART + TITLE + DESCRIPTION } input GalleryAlbumCreateInput {