From 91c04b5ed96797a7bd95b5f644fa4a1a93e3a3f7 Mon Sep 17 00:00:00 2001 From: Ranjan Shrestha Date: Mon, 6 Jul 2026 12:25:29 +0545 Subject: [PATCH 1/2] fix: update model configs; add output schema; add retry mechanism --- apps/reports/ai_features/extraction.py | 100 +++++++++++++++++++++---- apps/reports/ai_features/llms.py | 7 +- apps/reports/ai_features/prompts.py | 52 +++++++++++-- 3 files changed, 136 insertions(+), 23 deletions(-) diff --git a/apps/reports/ai_features/extraction.py b/apps/reports/ai_features/extraction.py index 8af3749..96c4f35 100644 --- a/apps/reports/ai_features/extraction.py +++ b/apps/reports/ai_features/extraction.py @@ -1,14 +1,17 @@ import base64 +import io import json import logging +import time from dataclasses import dataclass, field import fitz from langchain_ollama import ChatOllama, OllamaEmbeddings from langchain_openai import ChatOpenAI +from PIL import Image from apps.reports.ai_features.llms import OllamaHandler -from apps.reports.ai_features.prompts import get_doc_summary_prompt +from apps.reports.ai_features.prompts import DOC_SUMMARY_SCHEMA, PAGE_SCHEMA, get_doc_summary_prompt from apps.reports.models import DocumentExtraction, DocumentExtractionStatus, Report logger = logging.getLogger(__name__) @@ -58,16 +61,70 @@ def handle_meta_info(self): class PdfExtraction(BaseExtraction): data: bytes + # Qwen2.5-VL uses a native dynamic-resolution vision encoder, so the number of + # vision tokens (and Ollama's memory/compute use) scales with input pixel count. + # Cap the longest side so oversized source pages can't blow up memory regardless + # of the render zoom. + MAX_IMAGE_DIMENSION = 1024 + + # Retries for a single page's LLM extraction call, covering transient network/ + # timeout errors as well as malformed or truncated JSON in the model's response. + MAX_PAGE_ATTEMPTS = 3 + PAGE_RETRY_DELAY_SECONDS = 3 + def img_to_base64(self, data: fitz.Pixmap) -> str: img_bytes = data.tobytes("png") + + image = Image.open(io.BytesIO(img_bytes)) + if max(image.size) > self.MAX_IMAGE_DIMENSION: + scale = self.MAX_IMAGE_DIMENSION / max(image.size) + new_size = (round(image.width * scale), round(image.height * scale)) + image = image.resize(new_size, Image.Resampling.LANCZOS) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + img_bytes = buffer.getvalue() + return base64.b64encode(img_bytes).decode("utf-8") - def pdf_to_images(self, zoom: float = 2.0): + def extract_page(self, page_idx: int, img_b64: str) -> dict | None: + """Run the LLM extraction call for a single page, retrying on failure.""" + message = self.llm_handler.construct_extraction_message(img_b64=img_b64) + + for attempt in range(1, self.MAX_PAGE_ATTEMPTS + 1): + try: + response = self.llm_chat_model.invoke([message], format=PAGE_SCHEMA) + if not isinstance(response.content, str): + raise TypeError("Response content is not a string") # noqa: TRY301 + return json.loads(response.content) + except Exception: + logger.warning( + "Page %s extraction attempt %s/%s failed", + page_idx + 1, + attempt, + self.MAX_PAGE_ATTEMPTS, + exc_info=True, + ) + if attempt < self.MAX_PAGE_ATTEMPTS: + time.sleep(self.PAGE_RETRY_DELAY_SECONDS) + + logger.error("Page %s extraction failed after %s attempts", page_idx + 1, self.MAX_PAGE_ATTEMPTS) + return None + + def pdf_to_images(self, zoom: float = 1.1): page_summaries = [] doc = fitz.open(stream=self.data, filetype="pdf") # Get the title and description self.handle_meta_info() + # Doc Summary In Pending State + doc_summary_obj = DocumentExtraction.objects.create( + report=self.report, + status=DocumentExtractionStatus.PENDING, + text="", + page_number=None, + chunk_type=DocumentExtraction.ExtractionType.DOCUMENT_SUMMARY, + embedding=None, + ) for page_idx in range(len(doc)): logger.info("Processing Page %s", page_idx + 1) @@ -76,12 +133,15 @@ def pdf_to_images(self, zoom: float = 2.0): pic = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False) img_b64 = self.img_to_base64(data=pic) - message = self.llm_handler.construct_extraction_message(img_b64=img_b64) - - response = self.llm_chat_model.invoke([message]) - if not isinstance(response.content, str): + result = self.extract_page(page_idx=page_idx, img_b64=img_b64) + if result is None: + DocumentExtraction.objects.create( + report=self.report, + status=DocumentExtractionStatus.FAILURE, + page_number=page_idx + 1, + chunk_type=DocumentExtraction.ExtractionType.EXTRACTED_CONTENT, + ) continue - result = json.loads(response.content) if "summary" in result and result["summary"]: page_summaries.append(result["summary"]) @@ -124,21 +184,29 @@ def pdf_to_images(self, zoom: float = 2.0): ) doc_summary_prompt = get_doc_summary_prompt(page_summaries=page_summaries) - doc_summary = self.llm_chat_model.invoke(doc_summary_prompt) + # This prompt concatenates every page's summary, so its input size scales with + # page count. Override back up to the original context window rather than the + # smaller per-page default, since a lower window here can silently truncate + # earlier page summaries out of the final document summary. + doc_summary = self.llm_chat_model.invoke( + doc_summary_prompt, + format=DOC_SUMMARY_SCHEMA, + options={"num_ctx": 8192}, + ) if not isinstance(doc_summary.content, str): return - doc_summary_json = json.loads(doc_summary.content) - if doc_summary_json and "doc_summary" in doc_summary_json: - DocumentExtraction.objects.create( - report=self.report, + try: + doc_summary_json = json.loads(doc_summary.content) + DocumentExtraction.objects.filter(pk=doc_summary_obj.pk).update( status=DocumentExtractionStatus.SUCCESS, text=doc_summary_json["doc_summary"], - page_number=None, - chunk_type=DocumentExtraction.ExtractionType.DOCUMENT_SUMMARY, embedding=self.llm_embedding_model.embed_query(doc_summary_json["doc_summary"]), ) - else: - logger.warning("The key doc_summary is missing in the output") + except (ValueError, KeyError): + logger.warning("Either key doc_summary is missing or malformed json in the output.") + DocumentExtraction.objects.filter(pk=doc_summary_obj.pk).update( + status=DocumentExtractionStatus.FAILURE, + ) @dataclass diff --git a/apps/reports/ai_features/llms.py b/apps/reports/ai_features/llms.py index 3f3192d..cd00563 100644 --- a/apps/reports/ai_features/llms.py +++ b/apps/reports/ai_features/llms.py @@ -11,7 +11,7 @@ @dataclass class LLMHandler: - temperature: float = 0.2 + temperature: float = 0.0 def construct_extraction_message(self, img_b64: str) -> HumanMessage: return HumanMessage( @@ -49,7 +49,10 @@ def load_chat_model(self): model=settings.LLM_MODEL_NAME, base_url=settings.LLM_OLLAMA_BASE_URL, temperature=self.temperature, - format="json", + # Sized for a single page (prompt + one page's vision tokens + JSON output). + # The multi-page document summary call needs a larger window and overrides + # this via `options={"num_ctx": ...}` at call time. + num_ctx=4096, client_kwargs={ "timeout": httpx.Timeout( connect=30.0, diff --git a/apps/reports/ai_features/prompts.py b/apps/reports/ai_features/prompts.py index a51225b..5650788 100644 --- a/apps/reports/ai_features/prompts.py +++ b/apps/reports/ai_features/prompts.py @@ -1,8 +1,50 @@ +PAGE_SCHEMA = { + "type": "object", + "properties": { + "extracted_text": {"type": "string"}, + "key_findings": {"type": "string"}, + "tables": { + "type": "array", + "items": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "headers": {"type": "array", "items": {"type": "string"}}, + "rows": {"type": "array", "items": {"type": "array", "items": {"type": "string"}}}, + }, + "required": ["title", "headers", "rows"], + }, + }, + "charts": { + "type": "array", + "items": { + "type": "object", + "properties": { + "type": {"type": "string"}, + "title": {"type": "string"}, + "description": {"type": "string"}, + }, + "required": ["type", "title", "description"], + }, + }, + "summary": {"type": "string"}, + }, + "required": ["extracted_text", "key_findings", "tables", "charts", "summary"], +} + +DOC_SUMMARY_SCHEMA = { + "type": "object", + "properties": { + "doc_summary": {"type": "string"}, + }, +} + PAGE_PROMPT = """ -You are analyzing an image of a document page. -Extract all content and return ONLY a valid JSON object with no explanation, no markdown, no backticks. +You are analyzing an image(scan) of a document page. +Extract all content and return ONLY a valid JSON object based on page schema +with no explanation, no markdown, no backticks. -Use exactly this structure: +Use this exact structure: { "extracted_text": "extracted texts of the page" "key_findings": "important phrases separated by a period", @@ -20,11 +62,11 @@ "description": "extract key information and describe what the chart shows" } ], - "summary": "brief summary of the page content including key information from tables and charts under 200 words.", + "summary": "brief summary of the extracted texts including key information from tables and charts under 200 words.", } Rules: -- Return ONLY the JSON object, nothing else +- Return ONLY the valid JSON object following the schema - If no tables found, return "tables": [] - If no charts found, return "charts": [] - If no key findings, return "key_findings": "" else express in phrases From 79cdc1bd8cb05df342c76cbe967835804c2305a7 Mon Sep 17 00:00:00 2001 From: Ranjan Shrestha Date: Mon, 6 Jul 2026 15:15:42 +0545 Subject: [PATCH 2/2] fix: add the filter for status --- apps/reports/ai_features/search_docs.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/apps/reports/ai_features/search_docs.py b/apps/reports/ai_features/search_docs.py index f492e76..0778609 100644 --- a/apps/reports/ai_features/search_docs.py +++ b/apps/reports/ai_features/search_docs.py @@ -7,7 +7,7 @@ from pgvector.django import CosineDistance from apps.reports.ai_features.llms import OllamaHandler -from apps.reports.models import DocumentExtraction, Report +from apps.reports.models import DocumentExtraction, DocumentExtractionStatus, Report class ChunkScore(typing.TypedDict): @@ -47,7 +47,9 @@ def get_scores(self, k_top: int = 100) -> QuerySet[DocumentExtraction]: """Calculate the cosine similarity score of each of the chunks.""" query_vector = self.generate_query_embedding() return ( - DocumentExtraction.objects.select_related("report") + DocumentExtraction.objects.filter(status=DocumentExtractionStatus.SUCCESS) + .filter(embedding__isnull=False) + .select_related("report") .annotate( score=1 - CosineDistance( @@ -103,4 +105,4 @@ def rank_reports(self, top_k: int = 10) -> list[Report]: ) } # Return ordered reports - return [reports_by_id[rep_id] for rep_id, _ in ranked_docs] + return [reports_by_id[rep_id] for rep_id, _ in ranked_docs[:top_k]]