refactor(ocr): Extract OCR functionality into a separate module

- Moved OCR-related methods from `mistralai_engine.py` to a new module `mistralai_ocr.py` to enhance code modularity and maintainability.
- Updated `handle_ocr` in `mistralai_engine.py` to utilize the `handle_ocr` function from the new module.
- Removed unused imports and corrected import paths in `mistralai_engine.py`.
- Fixed a typo in a model description string.
- Cleaned up docstring comments in `conversation_tab.py` for better clarity.

Author: AAClause

basilisk/gui/conversation_tab.py

@@ -1172,7 +1172,6 @@ def _check_attachments_valid(self) -> bool:
 		return not invalid_found
 
 	def create_ocr_widget(self):
-		"""Create the OCR widget for the conversation tab."""
 		"""Create and configure the OCR button"""
 		self.ocr_button = wx.Button(
 			self,

basilisk/provider_engine/mistralai_engine.py

@@ -7,30 +7,22 @@
 from __future__ import annotations
 
 import logging
-from datetime import datetime
 from functools import cached_property
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator
 
 from mistralai import Mistral
-from mistralai.models import (
-	ChatCompletionResponse,
-	CompletionEvent,
-	OCRResponse,
-)
+from mistralai.models import ChatCompletionResponse, CompletionEvent
 from mistralai.utils.eventstreaming import EventStream
-from platformdirs import user_documents_dir
 
 from basilisk.conversation import (
 	Conversation,
 	Message,
 	MessageBlock,
 	MessageRoleEnum,
 )
-from basilisk.conversation.attached_file import (
-	AttachmentFile,
-	AttachmentFileTypes,
-)
+from basilisk.conversation.attached_file import AttachmentFile
+
+from .mistralai_ocr import handle_ocr
 
 if TYPE_CHECKING:
 	from multiprocessing import Queue
@@ -48,14 +40,14 @@ class MistralAIEngine(BaseEngine):
 	image and document capabilities.
 
 	Attributes:
-		capabilities: Set of supported capabilities including text, image, STT, and TTS.
+		capabilities: Set of supported capabilities including text, image, document, and OCR.
 	"""
 
 	capabilities: set[ProviderCapability] = {
+		ProviderCapability.TEXT,
 		ProviderCapability.IMAGE,
 		ProviderCapability.DOCUMENT,
 		ProviderCapability.OCR,
-		ProviderCapability.TEXT,
 	}
 	supported_attachment_formats: set[str] = {
 		"image/gif",
@@ -102,7 +94,7 @@ def models(self) -> list[ProviderAIModel]:
 				id="ministral-3b-latest",
 				name="Ministral 3B",
 				# Translators: This is a model description
-				description=_("World’s best edge model"),
+				description=_("World's best edge model"),
 				context_window=131000,
 				max_temperature=1.0,
 				default_temperature=0.7,
@@ -311,235 +303,6 @@ def completion_response_without_stream(
 		)
 		return new_block
 
-	@staticmethod
-	def _ocr_upload(client: Mistral, file: dict[str, Any]) -> str:
-		"""Uploads a file for OCR processing.
-
-		Args:
-			client: Mistral client instance.
-			file: The file to upload.
-
-		Returns:
-			The signed URL for the uploaded file.
-		"""
-		uploaded_pdf = client.files.upload(file=file, purpose="ocr")
-		return client.files.get_signed_url(file_id=uploaded_pdf.id).url
-
-	@staticmethod
-	def _ocr_process(
-		client: Mistral,
-		document: dict[str, Any],
-		include_image_base64: bool = True,
-		**kwargs,
-	) -> OCRResponse:
-		"""Processes a document for OCR.
-
-		Args:
-			client: Mistral client instance
-			document: The document to process.
-			include_image_base64: Whether to include image base64.
-			**kwargs: Additional keyword arguments.
-
-		Returns:
-			The OCR response.
-		"""
-		return client.ocr.process(
-			model="mistral-ocr-latest",
-			document=document,
-			include_image_base64=include_image_base64,
-			**kwargs,
-		)
-
-	@staticmethod
-	def _ocr_result(ocr_response: OCRResponse, file_path: str) -> bool:
-		"""Extracts text from the OCR response.
-
-		Args:
-			ocr_response: The OCR response.
-			file_path: The file path to save the extracted text.
-
-		Returns:
-			True if text was saved to a file, False otherwise.
-		"""
-		if not ocr_response:
-			return False
-
-		parent_dir = Path(file_path).parent
-		parent_dir.mkdir(exist_ok=True, parents=True)
-
-		with open(file_path, "w", encoding="UTF-8") as file:
-			for page_number, page in enumerate(ocr_response.pages, start=1):
-				file.write(page.markdown)
-				file.write(f"\n\n_----------_{page_number + 1}\n\n")
-		return True
-
-	@staticmethod
-	def _update_ocr_progress(message, progress=None, result_queue=None):
-		"""Updates OCR processing progress.
-
-		Args:
-			message: Progress message to log
-			progress: Optional numeric progress value
-			result_queue: Queue to send progress messages
-		"""
-		import sys
-
-		sys.stdout.write(f"Progress: {message}\n")
-		if result_queue:
-			try:
-				if message:
-					result_queue.put(("message", message))
-				if progress is not None:
-					result_queue.put(("progress", progress))
-			except Exception as e:
-				sys.stderr.write(f"Error sending progress update: {str(e)}\n")
-
-	@staticmethod
-	def _process_url_attachment(client, attachment, result_queue=None) -> str:
-		"""Process a URL-based attachment for OCR.
-
-		Args:
-			client: Mistral client instance
-			attachment: The attachment to process
-			result_queue: Queue for progress updates
-
-		Returns:
-			Path to the output file if successful, empty string otherwise
-		"""
-		from pathlib import Path
-
-		path = Path(user_documents_dir()) / "basilisk_ocr"
-		path.mkdir(exist_ok=True, parents=True)
-		output_file = (
-			path / f"{datetime.now().isoformat().replace(':', '-')}.md"
-		)
-
-		MistralAIEngine._update_ocr_progress(
-			f"Processing OCR for URL: {attachment.location}",
-			result_queue=result_queue,
-		)
-
-		result = MistralAIEngine._ocr_result(
-			MistralAIEngine._ocr_process(
-				client=client,
-				document={
-					"type": "document_url",
-					"document_url": attachment.url,
-				},
-				include_image_base64=True,
-			),
-			file_path=str(output_file),
-		)
-
-		return str(output_file) if result else ""
-
-	@staticmethod
-	def _process_file_attachment(client, attachment, result_queue=None) -> str:
-		"""Process a file-based attachment for OCR.
-
-		Args:
-			client: Mistral client instance
-			attachment: The attachment to process
-			result_queue: Queue for progress updates
-
-		Returns:
-			Path to the output file if successful, empty string otherwise
-		"""
-		import os
-		import sys
-		from pathlib import Path
-
-		output_file = Path(attachment.location).with_suffix(".md")
-		MistralAIEngine._update_ocr_progress(
-			f"Processing OCR for file: {attachment.name}",
-			result_queue=result_queue,
-		)
-
-		# Check if the file exists and is readable
-		if not os.path.exists(attachment.location):
-			MistralAIEngine._update_ocr_progress(
-				f"Warning: File not found: {attachment.location}",
-				result_queue=result_queue,
-			)
-			return ""
-
-		# Upload the file for OCR processing
-		with open(attachment.location, "rb") as f:
-			signed_url = MistralAIEngine._ocr_upload(
-				client=client, file={"file_name": attachment.name, "content": f}
-			)
-			sys.stdout.write(f"Signed URL: {signed_url}\n")
-
-		# Process the uploaded file
-		result = MistralAIEngine._ocr_result(
-			MistralAIEngine._ocr_process(
-				client=client,
-				document={"type": "document_url", "document_url": signed_url},
-				include_image_base64=True,
-			),
-			file_path=str(output_file),
-		)
-
-		return str(output_file) if result else ""
-
-	@staticmethod
-	def _process_single_attachment(
-		client, attachment, index, total, result_queue=None
-	) -> str:
-		"""Process a single attachment for OCR.
-
-		Args:
-			client: Mistral client instance
-			attachment: The attachment to process
-			index: Current attachment index
-			total: Total number of attachments
-			result_queue: Queue for progress updates
-
-		Returns:
-			Path to the output file if successful, empty string otherwise
-		"""
-		import traceback
-
-		MistralAIEngine._update_ocr_progress(
-			f"Processing attachment {index + 1}/{total}: {attachment.name}",
-			progress=int(100 * index / total),
-			result_queue=result_queue,
-		)
-
-		try:
-			if attachment.type == AttachmentFileTypes.URL:
-				output_file = MistralAIEngine._process_url_attachment(
-					client, attachment, result_queue
-				)
-			else:
-				output_file = MistralAIEngine._process_file_attachment(
-					client, attachment, result_queue
-				)
-
-			if output_file:
-				MistralAIEngine._update_ocr_progress(
-					f"OCR completed for {attachment.name}. Saved to {output_file}",
-					result_queue=result_queue,
-				)
-				return output_file
-			else:
-				MistralAIEngine._update_ocr_progress(
-					f"OCR failed for {attachment.name}. No text extracted.",
-					result_queue=result_queue,
-				)
-				return ""
-
-		except Exception as e:
-			error_trace = traceback.format_exc()
-			MistralAIEngine._update_ocr_progress(
-				f"Error processing attachment {attachment.name}: {str(e)}",
-				result_queue=result_queue,
-			)
-			import sys
-
-			sys.stderr.write(f"OCR error: {str(e)}\n{error_trace}\n")
-			return ""
-
 	@staticmethod
 	def handle_ocr(
 		api_key: str,
@@ -560,44 +323,6 @@ def handle_ocr(
 		Returns:
 			List of file paths containing the extracted text.
 		"""
-		import sys
-		import traceback
-
-		from mistralai import Mistral
-
-		try:
-			# Create a new client in the subprocess
-			MistralAIEngine._update_ocr_progress(
-				"Initializing OCR processor...", result_queue=result_queue
-			)
-			client = Mistral(api_key=api_key, server_url=base_url)
-
-			if not attachments:
-				return "error", "No attachments for OCR processing"
-
-			total_attachments = len(attachments)
-			output_files = []
-
-			for i, attachment in enumerate(attachments):
-				# Check for cancellation
-				if cancel_flag and cancel_flag.value:
-					return "result", output_files
-
-				output_file = MistralAIEngine._process_single_attachment(
-					client, attachment, i, total_attachments, result_queue
-				)
-				if output_file:
-					output_files.append(output_file)
-
-			# Final update
-			MistralAIEngine._update_ocr_progress(
-				f"OCR completed for {len(output_files)} of {total_attachments} attachments",
-				progress=100,
-				result_queue=result_queue,
-			)
-			return output_files
-
-		except Exception as e:
-			error_trace = traceback.format_exc()
-			sys.stderr.write(f"OCR process error: {str(e)}\n{error_trace}\n")
-			return "error", f"OCR process error: {str(e)}"
+		return handle_ocr(
+			api_key, base_url, attachments, cancel_flag, result_queue
+		)