diff --git a/mxtoai/agents/email_agent.py b/mxtoai/agents/email_agent.py index d3b330d..6815883 100644 --- a/mxtoai/agents/email_agent.py +++ b/mxtoai/agents/email_agent.py @@ -53,7 +53,7 @@ # Import citation management and web search tools from mxtoai.scripts.report_formatter import ReportFormatter -from mxtoai.scripts.visual_qa import azure_visualizer +from mxtoai.scripts.visual_qa import azure_visualizer_from_content from mxtoai.tools.attachment_processing_tool import AttachmentProcessingTool from mxtoai.tools.citation_aware_visit_tool import CitationAwareVisitTool from mxtoai.tools.deep_research_tool import DeepResearchTool @@ -94,7 +94,12 @@ class EmailAgent: """ def __init__( - self, email_request: EmailRequest, attachment_dir: str = "email_attachments", verbose: bool = False, enable_deep_research: bool = False, attachment_info: list[dict] | None = None + self, + email_request: EmailRequest, + attachment_dir: str = "email_attachments", + verbose: bool = False, + enable_deep_research: bool = False, + attachment_info: list[dict] | None = None, ): """ Initialize the email agent with tools for different operations. @@ -150,7 +155,7 @@ def __init__( self.scheduled_tasks_tool, self.delete_scheduled_tasks_tool, self.references_generator_tool, - azure_visualizer, + azure_visualizer_from_content, ] # Add all available search tools @@ -485,7 +490,9 @@ def _create_task_template( f"Process this email according to the '{handle}' instruction type.\n", email_context, distilled_section, - RESEARCH_GUIDELINES["mandatory"] if deep_research_mandatory and self.enable_deep_research else RESEARCH_GUIDELINES["optional"], + RESEARCH_GUIDELINES["mandatory"] + if deep_research_mandatory and self.enable_deep_research + else RESEARCH_GUIDELINES["optional"], attachment_task, handle_specific_template, output_template, @@ -666,7 +673,11 @@ def _process_agent_result( logger.info(f"Scheduled task created successfully with ID: {tool_output['task_id']}") else: error_msg = tool_output.get("message", "Scheduled task creation failed") - error_type = "Scheduled Task Limit Exceeded" if tool_output.get("error") == "Task limit exceeded" else "Scheduled Task Error" + error_type = ( + "Scheduled Task Limit Exceeded" + if tool_output.get("error") == "Task limit exceeded" + else "Scheduled Task Error" + ) errors_list.append(ProcessingError(message=error_type, details=error_msg)) if tool_output.get("error") == "Task limit exceeded": logger.warning(f"Scheduled task limit exceeded: {error_msg}") @@ -939,7 +950,9 @@ def limited_forward(*args, **kwargs): current_active_count = count_active_tasks_for_user(session, user_email) if current_active_count >= max_calls: - logger.warning(f"Scheduled task limit reached ({max_calls} active tasks per email). User has {current_active_count} active tasks.") + logger.warning( + f"Scheduled task limit reached ({max_calls} active tasks per email). User has {current_active_count} active tasks." + ) return { "success": False, "error": "Task limit exceeded", @@ -954,7 +967,6 @@ def limited_forward(*args, **kwargs): # Call the original method return original_forward(*args, **kwargs) - # Replace the forward method base_tool.forward = limited_forward @@ -974,9 +986,12 @@ def _finalize_response_with_citations(self, content: str) -> str: if self.context.has_citations(): # Check if content already contains a References or Sources section to avoid duplication import re + existing_references_pattern = r"(^|\n)#{1,3}\s*(References|Sources|Bibliography)\s*$" if re.search(existing_references_pattern, content, re.MULTILINE | re.IGNORECASE): - logger.warning("Content already contains a References/Sources section - skipping automatic references to avoid duplication") + logger.warning( + "Content already contains a References/Sources section - skipping automatic references to avoid duplication" + ) return content references_section = self.context.get_references_section() diff --git a/mxtoai/request_context.py b/mxtoai/request_context.py index dc18c97..61bb747 100644 --- a/mxtoai/request_context.py +++ b/mxtoai/request_context.py @@ -64,7 +64,7 @@ def load_attachment(self, filename: str, file_path: str, content_type: str, size "filename": filename, "contentType": content_type, "size": size, - "original_path": file_path + "original_path": file_path, } logger.debug(f"Loaded attachment into memory: {filename} ({size} bytes)") return True @@ -90,7 +90,7 @@ def has_attachment(self, filename: str) -> bool: class CitationManager: - """Per-request citation manager without threading complexity.""" + """Per-request citation manager.""" def __init__(self): self._citations = CitationCollection() @@ -125,7 +125,7 @@ def add_web_source(self, url: str, title: str, description: str | None = None, * url=url, date_accessed=datetime.now(timezone.utc).strftime("%Y-%m-%d"), source_type="web", - description=description + description=description, ) self._citations.add_source(source) @@ -150,7 +150,7 @@ def add_attachment_source(self, filename: str, description: str | None = None) - filename=filename, date_accessed=datetime.now(timezone.utc).strftime("%Y-%m-%d"), source_type="attachment", - description=description or "processed attachment" + description=description or "processed attachment", ) self._citations.add_source(source) @@ -170,7 +170,7 @@ def add_api_source(self, title: str, description: str | None = None) -> str: title=sanitized_title, date_accessed=datetime.now(timezone.utc).strftime("%Y-%m-%d"), source_type="api", - description=description or "API data" + description=description or "API data", ) self._citations.add_source(source) @@ -179,8 +179,7 @@ def add_api_source(self, title: str, description: str | None = None) -> str: def get_citations(self) -> CitationCollection: """Get a copy of the current citations.""" return CitationCollection( - sources=self._citations.sources.copy(), - references_section=self._citations.references_section + sources=self._citations.sources.copy(), references_section=self._citations.references_section ) def generate_references_section(self) -> str: @@ -189,12 +188,7 @@ def generate_references_section(self) -> str: return "" # Categorize sources - source_categories = { - "visited": [], - "search": [], - "attachment": [], - "api": [] - } + source_categories = {"visited": [], "search": [], "attachment": [], "api": []} for source in self._citations.sources: if source.source_type == "web": @@ -213,7 +207,7 @@ def generate_references_section(self) -> str: ("visited", "#### Visited Pages", lambda s: f"{s.id}. [{s.title}]({s.url})"), ("search", "#### Search Results", lambda s: f"{s.id}. [{s.title}]({s.url})"), ("attachment", "#### Attachments", lambda s: f"{s.id}. {s.filename}"), - ("api", "#### Data Sources", lambda s: f"{s.id}. {s.title}") + ("api", "#### Data Sources", lambda s: f"{s.id}. {s.title}"), ] for category, header, formatter in section_configs: @@ -265,10 +259,7 @@ def _load_attachments(self, attachment_info: list[dict]): """Load attachments from disk into memory store.""" for info in attachment_info: success = self.attachment_service.load_attachment( - filename=info["filename"], - file_path=info["path"], - content_type=info["type"], - size=info["size"] + filename=info["filename"], file_path=info["path"], content_type=info["type"], size=info["size"] ) if success: logger.debug(f"Successfully loaded attachment: {info['filename']}") diff --git a/mxtoai/scripts/email_processor.py b/mxtoai/scripts/email_processor.py deleted file mode 100644 index 10e960a..0000000 --- a/mxtoai/scripts/email_processor.py +++ /dev/null @@ -1,152 +0,0 @@ -import os -from email import policy -from email.parser import BytesParser -from pathlib import Path -from typing import Any - - -class EmailProcessor: - """ - Process email content and attachments. - """ - - def __init__(self, temp_dir: str = "email_attachments"): - """ - Initialize the EmailProcessor. - - Args: - temp_dir: Directory to store extracted attachments - - """ - self.temp_dir = temp_dir - os.makedirs(temp_dir, exist_ok=True) - - def process_email_file(self, email_file: str) -> dict[str, Any]: - """ - Process an email file and extract its content and attachments. - - Args: - email_file: Path to the email file (EML format) - - Returns: - Dict containing email metadata, body, and attachment paths - - """ - with open(email_file, "rb") as fp: - msg = BytesParser(policy=policy.default).parse(fp) - - # Extract basic metadata - metadata = { - "subject": msg.get("subject", ""), - "from": msg.get("from", ""), - "to": msg.get("to", ""), - "date": msg.get("date", ""), - } - - body = self._extract_body(msg) - attachments = self._extract_attachments(msg, email_file) - research_instructions = self._extract_research_instructions(body) - - return { - "metadata": metadata, - "body": body, - "attachments": attachments, - "research_instructions": research_instructions, - "attachment_dir": os.path.join(self.temp_dir, Path(email_file).stem) if attachments else None, - } - - def _extract_body(self, msg) -> str: - """ - Extract the body content from the email message. - - Args: - msg: Email message object - - Returns: - Email body as plain text - - """ - body = "" - - # First, try to get the plain text body - if msg.is_multipart(): - for part in msg.iter_parts(): - content_type = part.get_content_type() - if content_type == "text/plain": - body = part.get_content() - break - if content_type == "text/html" and not body: - # Use HTML if plain text isn't available - html_body = part.get_content() - # Simple HTML to text conversion (can be improved) - body = self._html_to_text(html_body) - elif msg.get_content_type() == "text/plain": - body = msg.get_content() - elif msg.get_content_type() == "text/html": - html_body = msg.get_content() - body = self._html_to_text(html_body) - - return body - - def _html_to_text(self, html: str) -> str: - """ - Convert HTML to plain text. - - Args: - html: HTML content - - Returns: - Plain text version of the HTML - - """ - # Simple implementation - can be improved with BeautifulSoup - import re - - text = re.sub(r"<[^>]+>", " ", html) - return re.sub(r"\s+", " ", text).strip() - - def _extract_attachments(self, msg, email_file: str) -> list[str]: - """ - Extract attachments from the email message. - - Args: - msg: Email message object - email_file: Original email file path (used for naming) - - Returns: - List of paths to extracted attachments - - """ - attachments = [] - attachment_dir = os.path.join(self.temp_dir, Path(email_file).stem) - os.makedirs(attachment_dir, exist_ok=True) - - if msg.is_multipart(): - for _, part in enumerate(msg.iter_parts()): - filename = part.get_filename() - if filename: - # Clean the filename - filename = Path(filename).name - filepath = os.path.join(attachment_dir, filename) - with open(filepath, "wb") as fp: - fp.write(part.get_payload(decode=True)) - attachments.append(filepath) - - return attachments - - def _extract_research_instructions(self, body: str) -> str: - """ - Extract research instructions from the email body. - This can be enhanced with NLP to better identify the actual request. - - Args: - body: Email body text - - Returns: - Extracted research instructions - - """ - # For now, we'll use a simple approach: the entire body is the instruction - # In a more sophisticated implementation, this could use NLP to identify - # specific instructions or questions - return body diff --git a/mxtoai/scripts/mdconvert.py b/mxtoai/scripts/mdconvert.py index 868b3af..aa61c69 100644 --- a/mxtoai/scripts/mdconvert.py +++ b/mxtoai/scripts/mdconvert.py @@ -26,7 +26,6 @@ import pptx # File-format detection -import puremagic import pydub import requests import speech_recognition as sr @@ -119,41 +118,25 @@ def __init__(self, title: Union[str, None] = None, text_content: str = ""): class DocumentConverter: """Abstract superclass of all DocumentConverters.""" - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: + """Convert content from memory. Must be implemented by subclasses.""" raise NotImplementedError class PlainTextConverter(DocumentConverter): """Anything with content type text/plain""" - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: - # Guess the content type from any file extension that might be around + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: content_type, _ = mimetypes.guess_type("__placeholder" + kwargs.get("file_extension", "")) - # Only accept text files if content_type is None or not content_type.startswith("text/"): return None - # Try to detect if file is binary - try: - with open(local_path, "rb") as file: - # Read first chunk of the file - chunk = file.read(1024) - if b"\0" in chunk: # Binary file detection - return None - - # Try to decode as UTF-8 - try: - chunk.decode("utf-8") - except UnicodeDecodeError: - return None - except Exception: + if b"\0" in content[:1024]: return None - # If we got here, it's safe to read as text try: - with open(local_path, encoding="utf-8") as fh: - text_content = fh.read() + text_content = content.decode("utf-8") return DocumentConverterResult( title=None, text_content=text_content, @@ -165,14 +148,16 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert class HtmlConverter(DocumentConverter): """Anything with content type text/html""" - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: - # Bail if not html + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() not in [".html", ".htm"]: return None - with open(local_path, encoding="utf-8") as fh: - return self._convert(fh.read()) + try: + html_content = content.decode("utf-8") + return self._convert(html_content) + except UnicodeDecodeError: + return None def _convert(self, html_content: str) -> Union[None, DocumentConverterResult]: """Helper function that converts and HTML string.""" @@ -201,8 +186,7 @@ def _convert(self, html_content: str) -> Union[None, DocumentConverterResult]: class WikipediaConverter(DocumentConverter): """Handle Wikipedia pages separately, focusing only on the main document content.""" - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: - # Bail if not Wikipedia + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() not in [".html", ".htm"]: return None @@ -210,16 +194,15 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert if not re.search(r"^https?:\/\/[a-zA-Z]{2,3}\.wikipedia.org\/", url): return None - # Parse the file - soup = None - with open(local_path, encoding="utf-8") as fh: - soup = BeautifulSoup(fh.read(), "html.parser") + try: + html_content = content.decode("utf-8") + soup = BeautifulSoup(html_content, "html.parser") + except UnicodeDecodeError: + return None - # Remove javascript and style blocks for script in soup(["script", "style"]): script.extract() - # Print only the main content body_elm = soup.find("div", {"id": "mw-content-text"}) title_elm = soup.find("span", {"class": "mw-page-title-main"}) @@ -246,8 +229,7 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert class YouTubeConverter(DocumentConverter): """Handle YouTube specially, focusing on the video title, description, and transcript.""" - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: - # Bail if not YouTube + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() not in [".html", ".htm"]: return None @@ -255,12 +237,12 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert if not url.startswith("https://www.youtube.com/watch?"): return None - # Parse the file - soup = None - with open(local_path, encoding="utf-8") as fh: - soup = BeautifulSoup(fh.read(), "html.parser") + try: + html_content = content.decode("utf-8") + soup = BeautifulSoup(html_content, "html.parser") + except UnicodeDecodeError: + return None - # Read the meta tags assert soup.title is not None assert soup.title.string is not None metadata: dict[str, str] = {"title": soup.title.string} @@ -270,7 +252,6 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert metadata[meta[a]] = meta.get("content", "") break - # We can also try to read the full description. This is more prone to breaking, since it reaches into the page implementation try: for script in soup(["script"]): content = script.text @@ -287,7 +268,6 @@ def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConvert except Exception: pass - # Start preparing the page webpage_text = "# YouTube\n" title = self._get(metadata, ["title", "og:title", "name"]) # type: ignore @@ -367,16 +347,25 @@ class PdfConverter(DocumentConverter): Converts PDFs to Markdown. Most style information is ignored, so the results are essentially plain-text. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a PDF + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() != ".pdf": return None - return DocumentConverterResult( - title=None, - text_content=pdfminer.high_level.extract_text(local_path), - ) + try: + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + text_content = pdfminer.high_level.extract_text(temp_file_path) + os.unlink(temp_file_path) + + return DocumentConverterResult( + title=None, + text_content=text_content, + ) + except Exception: + return None class DocxConverter(HtmlConverter): @@ -384,17 +373,19 @@ class DocxConverter(HtmlConverter): Converts DOCX files to Markdown. Style information (e.g.m headings) and tables are preserved where possible. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a DOCX + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() != ".docx": return None - result = None - with open(local_path, "rb") as docx_file: - result = mammoth.convert_to_html(docx_file) + try: + from io import BytesIO + + result = mammoth.convert_to_html(BytesIO(content)) html_content = result.value return self._convert(html_content) + except Exception: + return None class XlsxConverter(HtmlConverter): @@ -402,23 +393,27 @@ class XlsxConverter(HtmlConverter): Converts XLSX files to Markdown, with each sheet presented as a separate Markdown table. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a XLSX + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() not in [".xlsx", ".xls"]: return None - sheets = pd.read_excel(local_path, sheet_name=None) - md_content = "" - for s in sheets: - md_content += f"## {s}\n" - html_content = sheets[s].to_html(index=False) - md_content += self._convert(html_content).text_content.strip() + "\n\n" + try: + from io import BytesIO - return DocumentConverterResult( - title=None, - text_content=md_content.strip(), - ) + sheets = pd.read_excel(BytesIO(content), sheet_name=None) + md_content = "" + for s in sheets: + md_content += f"## {s}\n" + html_content = sheets[s].to_html(index=False) + md_content += self._convert(html_content).text_content.strip() + "\n\n" + + return DocumentConverterResult( + title=None, + text_content=md_content.strip(), + ) + except Exception: + return None class PptxConverter(HtmlConverter): @@ -426,70 +421,69 @@ class PptxConverter(HtmlConverter): Converts PPTX files to Markdown. Supports heading, tables and images with alt text. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a PPTX + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() != ".pptx": return None - md_content = "" + try: + from io import BytesIO + + md_content = "" + + presentation = pptx.Presentation(BytesIO(content)) + slide_num = 0 + for slide in presentation.slides: + slide_num += 1 + + md_content += f"\n\n\n" + + title = slide.shapes.title + for shape in slide.shapes: + if self._is_picture(shape): + alt_text = "" + with contextlib.suppress(Exception): + alt_text = shape._element._nvXxPr.cNvPr.attrib.get("descr", "") + + filename = re.sub(r"\W", "", shape.name) + ".jpg" + md_content += "\n![" + (alt_text if alt_text else shape.name) + "](" + filename + ")\n" + + if self._is_table(shape): + html_table = "" + first_row = True + for row in shape.table.rows: + html_table += "" + for cell in row.cells: + if first_row: + html_table += "" + else: + html_table += "" + html_table += "" + first_row = False + html_table += "
" + html.escape(cell.text) + "" + html.escape(cell.text) + "
" + md_content += "\n" + self._convert(html_table).text_content.strip() + "\n" + + elif shape.has_text_frame: + if shape == title: + md_content += "# " + shape.text.lstrip() + "\n" + else: + md_content += shape.text + "\n" - presentation = pptx.Presentation(local_path) - slide_num = 0 - for slide in presentation.slides: - slide_num += 1 - - md_content += f"\n\n\n" - - title = slide.shapes.title - for shape in slide.shapes: - # Pictures - if self._is_picture(shape): - # https://github.com/scanny/python-pptx/pull/512#issuecomment-1713100069 - alt_text = "" - with contextlib.suppress(Exception): - alt_text = shape._element._nvXxPr.cNvPr.attrib.get("descr", "") - - # A placeholder name - filename = re.sub(r"\W", "", shape.name) + ".jpg" - md_content += "\n![" + (alt_text if alt_text else shape.name) + "](" + filename + ")\n" - - # Tables - if self._is_table(shape): - html_table = "" - first_row = True - for row in shape.table.rows: - html_table += "" - for cell in row.cells: - if first_row: - html_table += "" - else: - html_table += "" - html_table += "" - first_row = False - html_table += "
" + html.escape(cell.text) + "" + html.escape(cell.text) + "
" - md_content += "\n" + self._convert(html_table).text_content.strip() + "\n" - - # Text areas - elif shape.has_text_frame: - if shape == title: - md_content += "# " + shape.text.lstrip() + "\n" - else: - md_content += shape.text + "\n" - - md_content = md_content.strip() - - if slide.has_notes_slide: - md_content += "\n\n### Notes:\n" - notes_frame = slide.notes_slide.notes_text_frame - if notes_frame is not None: - md_content += notes_frame.text md_content = md_content.strip() - return DocumentConverterResult( - title=None, - text_content=md_content.strip(), - ) + if slide.has_notes_slide: + md_content += "\n\n### Notes:\n" + notes_frame = slide.notes_slide.notes_text_frame + if notes_frame is not None: + md_content += notes_frame.text + md_content = md_content.strip() + + return DocumentConverterResult( + title=None, + text_content=md_content.strip(), + ) + except Exception: + return None def _is_picture(self, shape): if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PICTURE: @@ -518,41 +512,32 @@ def _get_metadata(self, local_path): class WavConverter(MediaConverter): """ - Converts WAV files to markdown via extraction of metadata (if `exiftool` is installed), and speech transcription (if `speech_recognition` is installed). + Converts WAV files to markdown via extraction of metadata and speech transcription. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a XLSX + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() != ".wav": return None md_content = "" - # Add metadata - metadata = self._get_metadata(local_path) - if metadata: - for f in [ - "Title", - "Artist", - "Author", - "Band", - "Album", - "Genre", - "Track", - "DateTimeOriginal", - "CreateDate", - "Duration", - ]: - if f in metadata: - md_content += f"{f}: {metadata[f]}\n" - - # Transcribe try: - transcript = self._transcribe_audio(local_path) - md_content += "\n\n### Audio Transcript:\n" + ("[No speech detected]" if transcript == "" else transcript) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + try: + transcript = self._transcribe_audio(temp_file_path) + md_content += "\n\n### Audio Transcript:\n" + ( + "[No speech detected]" if transcript == "" else transcript + ) + except Exception: + md_content += "\n\n### Audio Transcript:\nError. Could not transcribe this audio." + + os.unlink(temp_file_path) except Exception: - md_content += "\n\n### Audio Transcript:\nError. Could not transcribe this audio." + return None return DocumentConverterResult( title=None, @@ -568,49 +553,29 @@ def _transcribe_audio(self, local_path) -> str: class Mp3Converter(WavConverter): """ - Converts MP3 and M4A files to markdown via extraction of metadata (if `exiftool` is installed), and speech transcription (if `speech_recognition` AND `pydub` are installed). + Converts MP3 and M4A files to markdown via extraction of metadata and speech transcription. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a MP3 + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() not in [".mp3", ".m4a"]: return None md_content = "" - # Add metadata - metadata = self._get_metadata(local_path) - if metadata: - for f in [ - "Title", - "Artist", - "Author", - "Band", - "Album", - "Genre", - "Track", - "DateTimeOriginal", - "CreateDate", - "Duration", - ]: - if f in metadata: - md_content += f"{f}: {metadata[f]}\n" - - # Transcribe handle, temp_path = tempfile.mkstemp(suffix=".wav") os.close(handle) try: + with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as temp_input: + temp_input.write(content) + temp_input_path = temp_input.name + if extension.lower() == ".mp3": - sound = pydub.AudioSegment.from_mp3(local_path) + sound = pydub.AudioSegment.from_mp3(temp_input_path) else: - sound = pydub.AudioSegment.from_file(local_path, format="m4a") + sound = pydub.AudioSegment.from_file(temp_input_path, format="m4a") sound.export(temp_path, format="wav") - _args = {} - _args.update(kwargs) - _args["file_extension"] = ".wav" - try: transcript = super()._transcribe_audio(temp_path).strip() md_content += "\n\n### Audio Transcript:\n" + ( @@ -619,10 +584,13 @@ def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: except Exception: md_content += "\n\n### Audio Transcript:\nError. Could not transcribe this audio." + os.unlink(temp_input_path) + except Exception: + return None finally: - os.unlink(temp_path) + with contextlib.suppress(OSError): + os.unlink(temp_path) - # Return the result return DocumentConverterResult( title=None, text_content=md_content.strip(), @@ -631,99 +599,75 @@ def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: class ZipConverter(DocumentConverter): """ - Extracts ZIP files to a permanent local directory and returns a listing of extracted files. + Extracts ZIP files and returns a listing of extracted files. """ def __init__(self, extract_dir: str = "downloads"): - """ - Initialize with path to extraction directory. - - Args: - extract_dir: The directory where files will be extracted. Defaults to "downloads" - - """ self.extract_dir = extract_dir - # Create the extraction directory if it doesn't exist os.makedirs(self.extract_dir, exist_ok=True) - def convert(self, local_path: str, **kwargs: Any) -> Union[None, DocumentConverterResult]: - # Bail if not a ZIP file + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") if extension.lower() != ".zip": return None - # Verify it's actually a ZIP file - if not zipfile.is_zipfile(local_path): - return None + try: + from io import BytesIO + + extracted_files = [] - # Extract all files and build list - extracted_files = [] - with zipfile.ZipFile(local_path, "r") as zip_ref: - # Extract all files - zip_ref.extractall(self.extract_dir) - # Get list of all files - for file_path in zip_ref.namelist(): - # Skip directories - if not file_path.endswith("/"): - extracted_files.append(self.extract_dir + "/" + file_path) + with zipfile.ZipFile(BytesIO(content), "r") as zip_ref: + zip_ref.extractall(self.extract_dir) + for file_path in zip_ref.namelist(): + if not file_path.endswith("/"): + extracted_files.append(self.extract_dir + "/" + file_path) - # Sort files for consistent output - extracted_files.sort() + extracted_files.sort() - # Build the markdown content - md_content = "Downloaded the following files:\n" - for file in extracted_files: - md_content += f"* {file}\n" + md_content = "Downloaded the following files:\n" + for file in extracted_files: + md_content += f"* {file}\n" - return DocumentConverterResult(title="Extracted Files", text_content=md_content.strip()) + return DocumentConverterResult(title="Extracted Files", text_content=md_content.strip()) + except Exception: + return None class ImageConverter(MediaConverter): """ - Converts images to markdown via extraction of metadata (if `exiftool` is installed), OCR (if `easyocr` is installed), and description via a multimodal LLM (if an mlm_client is configured). + Converts images to Markdown, with optional vision-language model descriptions. """ - def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: - # Bail if not a XLSX + def convert_content(self, content: bytes, **kwargs: Any) -> Union[None, DocumentConverterResult]: extension = kwargs.get("file_extension", "") - if extension.lower() not in [".jpg", ".jpeg", ".png"]: + if extension.lower() not in [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".webp"]: return None md_content = "" + filename = kwargs.get("filename", "image" + extension) - # Add metadata - metadata = self._get_metadata(local_path) - if metadata: - for f in [ - "ImageSize", - "Title", - "Caption", - "Description", - "Keywords", - "Artist", - "Author", - "DateTimeOriginal", - "CreateDate", - "GPSPosition", - ]: - if f in metadata: - md_content += f"{f}: {metadata[f]}\n" - - # Try describing the image with GPTV - mlm_client = kwargs.get("mlm_client") - mlm_model = kwargs.get("mlm_model") - if mlm_client is not None and mlm_model is not None: - md_content += ( - "\n# Description:\n" - + self._get_mlm_description( - local_path, extension, mlm_client, mlm_model, prompt=kwargs.get("mlm_prompt") - ).strip() - + "\n" - ) + description = None + client = kwargs.get("mlm_client") + model = kwargs.get("mlm_model") + if client and model: + try: + with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + description = self._get_mlm_description(temp_file_path, extension, client, model) + os.unlink(temp_file_path) + except Exception: + pass + + if description: + md_content += f"## Image Description\n{description}\n" + + md_content += f"\n![{filename}]({filename})\n" return DocumentConverterResult( - title=None, - text_content=md_content, + title=filename, + text_content=md_content.strip(), ) def _get_mlm_description(self, local_path, extension, client, model, prompt=None): @@ -769,8 +713,7 @@ class UnsupportedFormatException(Exception): class MarkdownConverter: """ - (In preview) An extremely simple text-based document reader, suitable for LLM use. - This reader will convert common file-types or webpages to Markdown. + Converts various file types to Markdown from memory content. """ def __init__( @@ -779,19 +722,11 @@ def __init__( mlm_client: Optional[Any] = None, mlm_model: Optional[Any] = None, ): - if requests_session is None: - self._requests_session = requests.Session() - else: - self._requests_session = requests_session - + self._requests_session = requests_session if requests_session is not None else requests.Session() self._mlm_client = mlm_client self._mlm_model = mlm_model - self._page_converters: list[DocumentConverter] = [] - # Register converters for successful browsing operations - # Later registrations are tried first / take higher priority than earlier registrations - # To this end, the most specific converters should appear below the most generic converters self.register_page_converter(PlainTextConverter()) self.register_page_converter(HtmlConverter()) self.register_page_converter(WikipediaConverter()) @@ -805,167 +740,54 @@ def __init__( self.register_page_converter(ZipConverter()) self.register_page_converter(PdfConverter()) - def convert( - self, source: Union[str, requests.Response], **kwargs: Any - ) -> DocumentConverterResult: # TODO: deal with kwargs - """ - Args: - - source: can be a string representing a path or url, or a requests.response object - - extension: specifies the file extension to use when interpreting the file. If None, infer from source (path, uri, content-type, etc.) - - """ - # Local path or url - if isinstance(source, str): - if source.startswith(("http://", "https://", "file://")): - return self.convert_url(source, **kwargs) - return self.convert_local(source, **kwargs) - # Request response - if isinstance(source, requests.Response): - return self.convert_response(source, **kwargs) - return None - - def convert_local(self, path: str, **kwargs: Any) -> DocumentConverterResult: # TODO: deal with kwargs - # Prepare a list of extensions to try (in order of priority) + def convert_content(self, content: bytes, **kwargs: Any) -> DocumentConverterResult: + """Convert content from memory bytes.""" ext = kwargs.get("file_extension") extensions = [ext] if ext is not None else [] - # Get extension alternatives from the path and puremagic - base, ext = os.path.splitext(path) - self._append_ext(extensions, ext) - self._append_ext(extensions, self._guess_ext_magic(path)) - - # Convert - return self._convert(path, extensions, **kwargs) - - # TODO what should stream's type be? - def convert_stream(self, stream: Any, **kwargs: Any) -> DocumentConverterResult: # TODO: deal with kwargs - # Prepare a list of extensions to try (in order of priority) - ext = kwargs.get("file_extension") - extensions = [ext] if ext is not None else [] + filename = kwargs.get("filename", "") + if filename: + base, file_ext = os.path.splitext(filename) + self._append_ext(extensions, file_ext) - # Save the file locally to a temporary file. It will be deleted before this method exits - handle, temp_path = tempfile.mkstemp() - fh = os.fdopen(handle, "wb") - result = None - try: - # Write to the temporary file - content = stream.read() - if isinstance(content, str): - fh.write(content.encode("utf-8")) - else: - fh.write(content) - fh.close() - - # Use puremagic to check for more extension options - self._append_ext(extensions, self._guess_ext_magic(temp_path)) - - # Convert - result = self._convert(temp_path, extensions, **kwargs) - # Clean up - finally: - with contextlib.suppress(Exception): - fh.close() - os.unlink(temp_path) - - return result - - def convert_url(self, url: str, **kwargs: Any) -> DocumentConverterResult: # TODO: fix kwargs type - # Send a HTTP request to the URL - user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0" - response = self._requests_session.get(url, stream=True, headers={"User-Agent": user_agent}) - response.raise_for_status() - return self.convert_response(response, **kwargs) - - def convert_response( - self, response: requests.Response, **kwargs: Any - ) -> DocumentConverterResult: # TODO fix kwargs type - # Prepare a list of extensions to try (in order of priority) - ext = kwargs.get("file_extension") - extensions = [ext] if ext is not None else [] - - # Guess from the mimetype - content_type = response.headers.get("content-type", "").split(";")[0] - self._append_ext(extensions, mimetypes.guess_extension(content_type)) - - # Read the content disposition if there is one - content_disposition = response.headers.get("content-disposition", "") - m = re.search(r"filename=([^;]+)", content_disposition) - if m: - base, ext = os.path.splitext(m.group(1).strip("\"'")) - self._append_ext(extensions, ext) - - # Read from the extension from the path - base, ext = os.path.splitext(urlparse(response.url).path) - self._append_ext(extensions, ext) - - # Save the file locally to a temporary file. It will be deleted before this method exits - handle, temp_path = tempfile.mkstemp() - fh = os.fdopen(handle, "wb") - result = None - try: - # Download the file - for chunk in response.iter_content(chunk_size=512): - fh.write(chunk) - fh.close() + return self._convert_content(content, extensions, **kwargs) - # Use puremagic to check for more extension options - self._append_ext(extensions, self._guess_ext_magic(temp_path)) - - # Convert - result = self._convert(temp_path, extensions, url=response.url) - except Exception: - pass - - # Clean up - finally: - with contextlib.suppress(Exception): - fh.close() - os.unlink(temp_path) - - return result - - def _convert(self, local_path: str, extensions: list[Union[str, None]], **kwargs) -> DocumentConverterResult: + def _convert_content(self, content: bytes, extensions: list[Union[str, None]], **kwargs) -> DocumentConverterResult: error_trace = "" - for ext in [*extensions, None]: # Try last with no extension + for ext in [*extensions, None]: for converter in self._page_converters: _kwargs = copy.deepcopy(kwargs) - # Overwrite file_extension appropriately if ext is None: if "file_extension" in _kwargs: del _kwargs["file_extension"] else: _kwargs.update({"file_extension": ext}) - # Copy any additional global options if "mlm_client" not in _kwargs and self._mlm_client is not None: _kwargs["mlm_client"] = self._mlm_client if "mlm_model" not in _kwargs and self._mlm_model is not None: _kwargs["mlm_model"] = self._mlm_model - # If we hit an error log it and keep trying res = None try: - res = converter.convert(local_path, **_kwargs) + res = converter.convert_content(content, **_kwargs) except Exception: error_trace = ("\n\n" + traceback.format_exc()).strip() if res is not None: - # Normalize the content res.text_content = "\n".join([line.rstrip() for line in re.split(r"\r?\n", res.text_content)]) res.text_content = re.sub(r"\n{3,}", "\n\n", res.text_content) - - # TODO return res - # If we got this far without success, report any exceptions if len(error_trace) > 0: - msg = f"Could not convert '{local_path}' to Markdown. File type was recognized as {extensions}. While converting the file, the following error was encountered:\n\n{error_trace}" + filename = kwargs.get("filename", "content") + msg = f"Could not convert '{filename}' to Markdown. File type was recognized as {extensions}. While converting the content, the following error was encountered:\n\n{error_trace}" raise FileConversionException(msg) - # Nothing can handle it! - msg = f"Could not convert '{local_path}' to Markdown. The formats {extensions} are not supported." + filename = kwargs.get("filename", "content") + msg = f"Could not convert '{filename}' to Markdown. The formats {extensions} are not supported." raise UnsupportedFormatException(msg) def _append_ext(self, extensions, ext): @@ -975,26 +797,7 @@ def _append_ext(self, extensions, ext): ext = ext.strip() if ext == "": return - # if ext not in extensions: - if True: - extensions.append(ext) - - def _guess_ext_magic(self, path): - """Use puremagic (a Python implementation of libmagic) to guess a file's extension based on the first few bytes.""" - # Use puremagic to guess - try: - guesses = puremagic.magic_file(path) - if len(guesses) > 0: - ext = guesses[0].extension.strip() - if len(ext) > 0: - return ext - except FileNotFoundError: - pass - except IsADirectoryError: - pass - except PermissionError: - pass - return None + extensions.append(ext) def register_page_converter(self, converter: DocumentConverter) -> None: """Register a page text converter.""" diff --git a/mxtoai/scripts/visual_qa.py b/mxtoai/scripts/visual_qa.py index e1f7542..5d14c34 100644 --- a/mxtoai/scripts/visual_qa.py +++ b/mxtoai/scripts/visual_qa.py @@ -1,8 +1,6 @@ import base64 import json -import mimetypes import os -import uuid from io import BytesIO from typing import Optional @@ -21,12 +19,12 @@ logger = get_logger("azure_visualizer") -def process_images_and_text(image_path: str, query: str, client: InferenceClient): +def process_images_and_text_from_content(content: bytes, query: str, client: InferenceClient): """ - Process images and text using the IDEFICS model. + Process images and text using the IDEFICS model from memory content. Args: - image_path: Path to the image file. + content: Image content as bytes. query: The question to ask about the image. client: Inference client for the model. @@ -45,22 +43,14 @@ def process_images_and_text(image_path: str, query: str, client: InferenceClient idefics_processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b-chatty") prompt_with_template = idefics_processor.apply_chat_template(messages, add_generation_prompt=True) - # load images from local directory - - # encode images to strings which can be sent to the endpoint - def encode_local_image(image_path): - # load image - image = Image.open(image_path).convert("RGB") - - # Convert the image to a base64 string + def encode_content_image(content: bytes): + image = Image.open(BytesIO(content)).convert("RGB") buffer = BytesIO() - image.save(buffer, format="JPEG") # Use the appropriate format (e.g., JPEG, PNG) + image.save(buffer, format="JPEG") base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8") - - # add string formatting required by the endpoint return f"data:image/jpeg;base64,{base64_image}" - image_string = encode_local_image(image_path) + image_string = encode_content_image(content) prompt_with_images = prompt_with_template.replace("", "![]({}) ").format(image_string) payload = { @@ -74,77 +64,52 @@ def encode_local_image(image_path): return json.loads(client.post(json=payload).decode())[0] -# Function to encode the image -def encode_image(image_path: str) -> str: +def encode_image_from_content(content: bytes) -> str: """ - Encode an image to base64 format. + Encode image content to base64 format. Args: - image_path: The path to the image file. + content: The image content as bytes. Returns: str: The base64 encoded string of the image. """ - if image_path.startswith("http"): - user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0" - request_kwargs = { - "headers": {"User-Agent": user_agent}, - "stream": True, - } - - # Send a HTTP request to the URL - response = requests.get(image_path, **request_kwargs) - response.raise_for_status() - content_type = response.headers.get("content-type", "") - - extension = mimetypes.guess_extension(content_type) - if extension is None: - extension = ".download" - - fname = str(uuid.uuid4()) + extension - download_path = os.path.abspath(os.path.join("downloads", fname)) - - with open(download_path, "wb") as fh: - for chunk in response.iter_content(chunk_size=512): - fh.write(chunk) - - image_path = download_path - - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode("utf-8") - - -headers = {"Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"} + return base64.b64encode(content).decode("utf-8") -def resize_image(image_path: str) -> str: +def resize_image_from_content(content: bytes) -> bytes: """ - Resize the image to half its original size. + Resize image content to half its original size. Args: - image_path: The path to the image file. + content: The image content as bytes. Returns: - str: The path to the resized image. + bytes: The resized image content as bytes. """ - img = Image.open(image_path) + img = Image.open(BytesIO(content)) width, height = img.size img = img.resize((int(width / 2), int(height / 2))) - directory = os.path.dirname(image_path) - filename = os.path.basename(image_path) - new_image_path = os.path.join(directory, f"resized_{filename}") - img.save(new_image_path) - return new_image_path + buffer = BytesIO() + img.save(buffer, format="JPEG") + return buffer.getvalue() + + +headers = {"Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"} class VisualQATool(Tool): name = "visualizer" - description = "A tool that can answer questions about attached images." + description = "A tool that can answer questions about attached images from memory content." inputs = { - "image_path": { - "description": "The path to the image on which to answer the question", + "content": { + "description": "The image content as bytes", + "type": "any", + }, + "mime_type": { + "description": "MIME type of the image", "type": "string", }, "question": {"description": "the question to answer", "type": "string", "nullable": True}, @@ -153,12 +118,13 @@ class VisualQATool(Tool): client = InferenceClient("HuggingFaceM4/idefics2-8b-chatty") - def forward(self, image_path: str, question: Optional[str] = None) -> str: + def forward(self, content: bytes, mime_type: str, question: Optional[str] = None) -> str: """ Process the image and return a short caption based on the content. Args: - image_path: The path to the image on which to answer the question. This should be a local path to downloaded image. + content: The image content as bytes. + mime_type: MIME type of the image. question: The question to answer. Returns: @@ -171,11 +137,11 @@ def forward(self, image_path: str, question: Optional[str] = None) -> str: add_note = True question = "Please write a detailed caption for this image." try: - output = process_images_and_text(image_path, question, self.client) + output = process_images_and_text_from_content(content, question, self.client) except Exception as e: if "Payload Too Large" in str(e): - new_image_path = resize_image(image_path) - output = process_images_and_text(new_image_path, question, self.client) + resized_content = resize_image_from_content(content) + output = process_images_and_text_from_content(resized_content, question, self.client) if add_note: output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}" @@ -184,12 +150,13 @@ def forward(self, image_path: str, question: Optional[str] = None) -> str: @tool -def visualizer(image_path: str, question: Optional[str] = None) -> str: +def visualizer_from_content(content: bytes, mime_type: str, question: Optional[str] = None) -> str: """ - A tool that can answer questions about attached images. + A tool that can answer questions about image content. Args: - image_path: The path to the image on which to answer the question. This should be a local path to downloaded image. + content: The image content as bytes. + mime_type: MIME type of the image. question: The question to answer. """ @@ -197,12 +164,11 @@ def visualizer(image_path: str, question: Optional[str] = None) -> str: if not question: add_note = True question = "Please write a detailed caption for this image." - if not isinstance(image_path, str): - msg = "You should provide at least `image_path` string argument to this tool!" - raise Exception(msg) - mime_type, _ = mimetypes.guess_type(image_path) - base64_image = encode_image(image_path) + if not mime_type: + mime_type = "image/jpeg" + + base64_image = encode_image_from_content(content) payload = { "model": "gpt-4o", @@ -231,12 +197,13 @@ def visualizer(image_path: str, question: Optional[str] = None) -> str: @tool -def azure_visualizer(image_path: str, question: Optional[str] = None) -> str: +def azure_visualizer_from_content(content: bytes, mime_type: str, question: Optional[str] = None) -> str: """ - A tool that can answer questions about attached images using Azure OpenAI. + A tool that can answer questions about image content using Azure OpenAI. Args: - image_path: The path to the image on which to answer the question. This should be a local path to downloaded image. + content: The image content as bytes. + mime_type: MIME type of the image. question: The question to answer. """ @@ -246,67 +213,85 @@ def azure_visualizer(image_path: str, question: Optional[str] = None) -> str: question = "Please write a detailed caption for this image." try: - # Get image MIME type - mime_type, _ = mimetypes.guess_type(image_path) if not mime_type: - mime_type = "image/jpeg" # Default to JPEG if can't determine + mime_type = "image/jpeg" - # Encode the image to base64 - base64_image = encode_image(image_path) + base64_image = encode_image_from_content(content) - # Format the content for the Azure OpenAI API - content = [ + content_payload = [ {"type": "text", "text": question}, {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}}, ] - # Get Azure OpenAI configuration model = f"azure/{os.getenv('MODEL_NAME', 'gpt-4o')}" api_key = os.getenv("MODEL_API_KEY") api_base = os.getenv("MODEL_ENDPOINT") api_version = os.getenv("MODEL_API_VERSION") - logger.info(f"Sending image to Azure OpenAI model: {model}") - # Call Azure OpenAI using litellm + logger.info(f"Sending image content to Azure OpenAI model: {model}") response = completion( model=model, - messages=[{"role": "user", "content": content}], + messages=[{"role": "user", "content": content_payload}], api_key=api_key, api_base=api_base, api_version=api_version, max_tokens=1000, ) - # Extract content from response if hasattr(response, "choices") and response.choices: if hasattr(response.choices[0], "message") and hasattr(response.choices[0].message, "content"): output = response.choices[0].message.content else: - # Fallback if direct access doesn't work output = response.choices[0].get("message", {}).get("content", "") else: - # Handle unusual response format output = str(response) if not output: msg = "Empty response from Azure OpenAI" raise Exception(msg) - # Add note if no question was provided if add_note: output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}" return output except Exception as e: - logger.error(f"Error in azure_visualizer: {e!s}") + logger.error(f"Error in azure_visualizer_from_content: {e!s}") - # Try resizing the image if it might be too large - if "too large" in str(e).lower() or "payload" in str(e).lower(): - try: - new_image_path = resize_image(image_path) - return azure_visualizer(new_image_path, question) - except Exception as resize_error: - logger.error(f"Error resizing image: {resize_error!s}") + try: + resized_content = resize_image_from_content(content) + base64_image = encode_image_from_content(resized_content) + + content_payload = [ + {"type": "text", "text": question}, + {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}}, + ] + + response = completion( + model=model, + messages=[{"role": "user", "content": content_payload}], + api_key=api_key, + api_base=api_base, + api_version=api_version, + max_tokens=1000, + ) + + if hasattr(response, "choices") and response.choices: + if hasattr(response.choices[0], "message") and hasattr(response.choices[0].message, "content"): + output = response.choices[0].message.content + else: + output = response.choices[0].get("message", {}).get("content", "") + else: + output = str(response) + + if add_note: + output = ( + f"You did not provide a particular question, so here is a detailed caption for the image: {output}" + ) - return f"Error processing image: {e!s}" + return output + + except Exception as retry_error: + logger.error(f"Error in azure_visualizer_from_content retry: {retry_error!s}") + msg = f"Failed to process image: {retry_error!s}" + raise Exception(msg) diff --git a/mxtoai/tools/__init__.py b/mxtoai/tools/__init__.py index d5d51d3..6996abe 100644 --- a/mxtoai/tools/__init__.py +++ b/mxtoai/tools/__init__.py @@ -1,9 +1,14 @@ # Tools package for email processing from mxtoai.tools.attachment_processing_tool import AttachmentProcessingTool +from mxtoai.tools.citation_aware_visit_tool import CitationAwareVisitTool from mxtoai.tools.deep_research_tool import DeepResearchTool from mxtoai.tools.delete_scheduled_tasks_tool import DeleteScheduledTasksTool +from mxtoai.tools.fallback_search_tool import FallbackWebSearchTool from mxtoai.tools.meeting_tool import MeetingTool +from mxtoai.tools.pdf_export_tool import PDFExportTool +from mxtoai.tools.references_generator_tool import ReferencesGeneratorTool from mxtoai.tools.scheduled_tasks_tool import ScheduledTasksTool +from mxtoai.tools.visual_qa_tool import VisualQATool # Web search tools from mxtoai.tools.web_search import BraveSearchTool, DDGSearchTool, GoogleSearchTool @@ -11,10 +16,15 @@ __all__ = [ "AttachmentProcessingTool", "BraveSearchTool", + "CitationAwareVisitTool", "DDGSearchTool", "DeepResearchTool", "DeleteScheduledTasksTool", + "FallbackWebSearchTool", "GoogleSearchTool", "MeetingTool", + "PDFExportTool", + "ReferencesGeneratorTool", "ScheduledTasksTool", + "VisualQATool", ] diff --git a/mxtoai/tools/attachment_processing_tool.py b/mxtoai/tools/attachment_processing_tool.py index 49dd1a2..3c18d75 100644 --- a/mxtoai/tools/attachment_processing_tool.py +++ b/mxtoai/tools/attachment_processing_tool.py @@ -1,384 +1,104 @@ -import json import os -import sys -from pathlib import Path -from typing import Any -from urllib.parse import unquote from smolagents import Tool -from smolagents.models import MessageRole, Model - -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import contextlib - -from scripts.mdconvert import MarkdownConverter from mxtoai._logging import get_logger -from mxtoai.request_context import RequestContext -from mxtoai.schemas import ToolOutputWithCitations +from mxtoai.scripts.mdconvert import MarkdownConverter -# Configure logger -logger = get_logger("attachment_tool") +logger = get_logger(__name__) class AttachmentProcessingTool(Tool): """ - Tool for processing various types of email attachments. - Handles documents using MarkdownConverter. For images, use the azure_visualizer tool directly. - """ - - name = "attachment_processor" - description = """Process and analyze email attachments to extract content and insights. - This tool can handle: - - Documents (PDFs, Office files, text files) - - Audio files (as transcripts) - - HTML files - - Markdown files - - NOTE: For image processing, please use the azure_visualizer tool directly. - This tool will skip image files and indicate they should be processed by azure_visualizer. - - The attachments parameter should be a list of dictionaries, where each dictionary contains: - - filename: Name of the file - - type: MIME type - - path: Full path to the file - - size: File size in bytes + Tool for processing email attachments from memory. """ + name = "attachment_processing" + description = "Process attachments from memory and convert them to text format" inputs = { - "attachments": { - "type": "array", - "description": "List of attachment dictionaries containing file information. Each dictionary must have 'filename', 'type', and 'size' keys. The 'path' key is optional - attachments will be processed from memory when available, falling back to file path if needed.", - "items": { - "type": "object", - "properties": { - "filename": {"type": "string", "description": "Name of the file"}, - "type": {"type": "string", "description": "MIME type or content type of the file"}, - "path": {"type": "string", "description": "Full path to the file (optional - used as fallback if memory content unavailable)"}, - "size": {"type": "integer", "description": "Size of the file in bytes"}, - }, - "required": ["filename", "type", "size"] - }, + "filename": { + "description": "Name of the attachment file to process", + "type": "string", }, - "mode": { + "type": { + "description": "MIME type of the attachment", "type": "string", - "description": "Processing mode: 'basic' for metadata only, 'full' for complete content analysis", - "enum": ["basic", "full"], - "default": "basic", - "nullable": True, + }, + "size": { + "description": "Size of the attachment in bytes", + "type": "integer", }, } - output_type = "object" + output_type = "string" - def __init__(self, context: RequestContext, model: Model | None = None, text_limit: int = 4000): - """ - Initialize the attachment processing tool. - - Args: - context: Request context containing email data and citation manager - model: Optional LLM model for content summarization - text_limit: Maximum text length to include in output - - """ + def __init__(self, context=None): super().__init__() - self.context = context - self.model = model - self.text_limit = text_limit self.converter = MarkdownConverter() + self.context = context - # Configure image extensions that should be handled by azure_visualizer - self.image_extensions = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".svg", ".tiff", ".ico"} - - logger.debug(f"AttachmentProcessingTool initialized with text_limit={text_limit}") - - # Set up attachments directory path - self.attachments_dir = Path( - os.path.abspath(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "attachments")) - ) - self.attachments_dir.mkdir(parents=True, exist_ok=True) - - def _validate_attachment_path(self, file_path: str) -> Path: + def forward(self, filename: str, type: str, size: int) -> str: """ - Validate and resolve the attachment file path. + Process attachment content from memory and return text representation. Args: - file_path: Path to the attachment file. + filename: Name of the attachment file + type: MIME type of the attachment + size: Size of the attachment in bytes Returns: - Path: The resolved file path. + str: Processed text content or error message """ try: - if not file_path: - msg = "Empty file path provided" - raise ValueError(msg) + if not self.context or not self.context.attachment_service: + return "Error: No attachment service available in request context" - # Clean up the path - file_path = file_path.strip("\"'") + if not self.context.attachment_service.has_attachment(filename): + return f"Error: Attachment '{filename}' not found" - # Try different path variations - paths_to_try = [ - Path(file_path), # Direct path - Path(unquote(file_path)), # URL decoded path - self.attachments_dir / Path(file_path).name, # Relative to attachments dir - ] + content = self.context.attachment_service.get_content(filename) + if content is None: + return f"Error: Could not retrieve content for '{filename}'" - for path in paths_to_try: - if path.is_file(): - return path.resolve() + processed_content = self._process_content_from_memory(content, filename, type) - paths_str = "\n- ".join(str(p) for p in paths_to_try) - msg = f"File not found at any of these locations:\n- {paths_str}" - raise FileNotFoundError(msg) + logger.info(f"Successfully processed attachment '{filename}' from memory") + return f"Processed attachment '{filename}':\n\n{processed_content}" except Exception as e: - logger.error(f"Error validating path {file_path}: {e!s}") - raise + logger.error(f"Error processing attachment '{filename}': {e!s}") + return f"Error processing attachment '{filename}': {e!s}" def _process_content_from_memory(self, content: bytes, filename: str, content_type: str) -> str: """ - Process document content from memory using MarkdownConverter. + Process file content directly from memory bytes. Args: - content: The file content as bytes + content: Raw file content as bytes filename: Name of the file for context - content_type: MIME type of the content + content_type: MIME type of the file Returns: - str: The text content extracted from the document. + str: Processed text content - """ - import tempfile + Raises: + ValueError: If processing fails + """ try: - # For text files, decode directly if content_type.startswith("text/"): try: return content.decode("utf-8") except UnicodeDecodeError: return content.decode("utf-8", errors="ignore") - # For other file types, create a temporary file for the converter - with tempfile.NamedTemporaryFile(suffix=f"_{filename}", delete=False) as temp_file: - temp_file.write(content) - temp_file_path = temp_file.name - - try: - result = self.converter.convert(temp_file_path) - if not result or not hasattr(result, "text_content"): - msg = f"Failed to convert document: {filename}" - raise ValueError(msg) - return result.text_content - finally: - # Clean up temporary file - import os - with contextlib.suppress(OSError): - os.unlink(temp_file_path) - - except Exception as e: - logger.error(f"Error converting document {filename} from memory: {e!s}") - raise - - def _process_document(self, file_path: Path) -> str: - """ - DEPRECATED: This method has been removed for security reasons. - Use _process_content_from_memory instead. - """ - logger.warning(f"_process_document is deprecated and removed for security. File path: {file_path}") - msg = "File path processing is deprecated for security. Use memory-based processing instead." - raise ValueError(msg) - - def forward(self, attachments: list[dict[str, Any]], mode: str = "basic") -> dict[str, Any]: - """ - Process email attachments synchronously with citation tracking. - - Args: - attachments: List of attachment dictionaries containing file information. - mode: Processing mode: 'basic' for metadata only, 'full' for complete content analysis. - - Returns: - str: JSON string of ToolOutputWithCitations containing processed attachments. - - """ - processed_attachments = [] - citation_ids = [] - - logger.info(f"Processing {len(attachments)} attachments in {mode} mode") - - for attachment in attachments: - try: - # Validate required fields - required_fields = ["filename", "type", "size"] - missing_fields = [field for field in required_fields if field not in attachment] - if missing_fields: - msg = f"Missing required fields in attachment: {missing_fields}" - raise ValueError(msg) - - filename = attachment["filename"] - content_type = attachment["type"] - logger.info(f"Processing attachment: {filename}") - - # Add citation for this attachment - citation_id = self.context.add_attachment_citation( - filename, - f"Email attachment ({content_type})" - ) - citation_ids.append(citation_id) - - # Skip image files - they should be handled by azure_visualizer directly - if content_type.startswith("image/"): - processed_attachments.append( - { - **attachment, - "citation_id": citation_id, - "content": { - "text": f"This is an image file that requires visual processing. [#{citation_id}]", - "type": "image", - "requires_visual_qa": True, - }, - } - ) - logger.info(f"Skipped image file: {filename} - use azure_visualizer tool instead") - continue - - # Try to get content from attachment service first - content = None - processing_source = "memory" - - if self.context.attachment_service.has_attachment(filename): - try: - file_content = self.context.attachment_service.get_content(filename) - if file_content: - content = self._process_content_from_memory(file_content, filename, content_type) - logger.debug(f"Processed {filename} from memory store") - else: - logger.warning(f"Attachment {filename} found in service but content is None") - except Exception as e: - logger.warning(f"Failed to process {filename} from memory: {e!s}, falling back to file path") - processing_source = "fallback" - - # Fall back to file path processing if memory processing failed or unavailable - if content is None and "path" in attachment: - logger.warning(f"File path processing is deprecated for security. Skipping {filename}. " - f"Use memory-based processing instead.") - processed_attachments.append({ - **attachment, - "citation_id": citation_id, - "error": "File path processing deprecated for security - use memory-based processing" - }) - continue - - # If we still don't have content, it's an error - if content is None: - error_msg = f"Could not process {filename}: no content available in memory or file path" - logger.error(error_msg) - processed_attachments.append({**attachment, "citation_id": citation_id, "error": error_msg}) - continue - - # If in full mode and model is available, generate a summary - summary = None - if mode == "full" and self.model and len(content) > 4000: - messages = [ - { - "role": MessageRole.SYSTEM, - "content": [ - { - "type": "text", - "text": f"Here is a file:\n### {filename}\n\n{content[: self.text_limit]}", - } - ], - }, - { - "role": MessageRole.USER, - "content": [ - { - "type": "text", - "text": "Please provide a comprehensive summary of this document in 5-7 sentences.", - } - ], - }, - ] - summary = self.model(messages).content - - processed_attachments.append( - { - **attachment, - "citation_id": citation_id, - "processing_source": processing_source, - "content": { - "text": f"{content[: self.text_limit] if len(content) > self.text_limit else content} [#{citation_id}]", - "type": "text", - "summary": summary, - }, - } - ) - logger.info(f"Successfully processed: {filename} (source: {processing_source})") + _, file_extension = os.path.splitext(filename) - except Exception as e: - logger.error(f"Error processing attachment {attachment.get('filename', 'unknown')}: {e!s}") - processed_attachments.append( - {**{k: v for k, v in attachment.items() if k in ["filename", "type", "size"]}, "citation_id": citation_id, "error": str(e)} - ) + result = self.converter.convert_content(content=content, filename=filename, file_extension=file_extension) - # Create structured output with citations - attachment_summary = self._create_attachment_summary(processed_attachments) - content = f"Processed {len(processed_attachments)} attachments:\n\n{attachment_summary}" + return result.text_content - result = ToolOutputWithCitations( - content=content, - metadata={ - "total_attachments": len(attachments), - "processed_successfully": len([a for a in processed_attachments if "error" not in a]), - "failed": len([a for a in processed_attachments if "error" in a]), - "citation_ids": citation_ids, - "attachments": processed_attachments, - } - ) - - return json.dumps(result.model_dump()) - - def _create_attachment_summary(self, attachments: list[dict[str, Any]]) -> str: - """ - Create a summary of processed attachments. - - Args: - attachments: List of processed attachment dictionaries. - - Returns: - str: Summary of processed attachments. - - """ - if not attachments: - return "No attachments processed." - - summary_parts = [] - successful = 0 - failed = 0 - images = 0 - - for att in attachments: - if "error" in att: - failed += 1 - summary_parts.append(f"Failed to process {att['filename']}: {att['error']}") - continue - - content = att.get("content", {}) - if content: - if content.get("type") == "image": - images += 1 - summary_parts.append(f"Image {att['filename']}: Requires visual processing") - elif content.get("type") == "text": - successful += 1 - summary_parts.append(f"Document: {att['filename']}") - if content.get("summary"): - summary_parts.append(f"Summary: {content['summary']}") - else: - text = content.get("text", "") - preview = text[:200] + "..." if len(text) > 200 else text - summary_parts.append(f"Preview: {preview}") - - status = f"Processed {successful} documents, {images} images pending visual processing" - if failed > 0: - status += f", {failed} failed" - - return status + "\n\n" + "\n\n".join(summary_parts) + except Exception as e: + logger.error(f"Error processing content from memory for {filename}: {e!s}") + msg = f"Failed to process content: {e!s}" + raise ValueError(msg) diff --git a/mxtoai/tools/visual_qa_tool.py b/mxtoai/tools/visual_qa_tool.py new file mode 100644 index 0000000..e36856e --- /dev/null +++ b/mxtoai/tools/visual_qa_tool.py @@ -0,0 +1,95 @@ +import mimetypes +from typing import Optional + +from smolagents import Tool + +from mxtoai._logging import get_logger +from mxtoai.scripts.visual_qa import azure_visualizer_from_content, visualizer_from_content + +logger = get_logger(__name__) + + +class VisualQATool(Tool): + """ + Tool for processing image attachments from memory using visual QA. + """ + + name = "visual_qa" + description = "Answer questions about image attachments from memory content" + inputs = { + "filename": { + "description": "Name of the image attachment to analyze", + "type": "string", + }, + "question": { + "description": "Question to ask about the image (optional - will provide caption if not specified)", + "type": "string", + "nullable": True, + }, + "use_azure": { + "description": "Whether to use Azure OpenAI (true) or standard OpenAI (false)", + "type": "boolean", + "nullable": True, + }, + } + output_type = "string" + + def __init__(self, context=None): + super().__init__() + self.context = context + + def forward(self, filename: str, question: Optional[str] = None, use_azure: Optional[bool] = None) -> str: + """ + Analyze an image attachment and answer questions about it. + + Args: + filename: Name of the image attachment + question: Question to ask about the image (optional) + use_azure: Whether to use Azure OpenAI (defaults to True) + + Returns: + str: Analysis result or error message + + """ + try: + if not self.context or not self.context.attachment_service: + return "Error: No attachment service available in request context" + + if not self.context.attachment_service.has_attachment(filename): + return f"Error: Attachment '{filename}' not found" + + content = self.context.attachment_service.get_content(filename) + if content is None: + return f"Error: Could not retrieve content for '{filename}'" + + metadata = self.context.attachment_service.get_metadata(filename) + if not metadata: + return f"Error: Could not retrieve metadata for '{filename}'" + + mime_type = metadata.get("type", "image/jpeg") + + # Check if it's an image file + if not mime_type.startswith("image/"): + guessed_type, _ = mimetypes.guess_type(filename) + if not guessed_type or not guessed_type.startswith("image/"): + return f"Error: '{filename}' is not an image file (type: {mime_type})" + mime_type = guessed_type + + # Use Azure by default + if use_azure is None: + use_azure = True + + # Process the image + if use_azure: + result = azure_visualizer_from_content(content, mime_type, question) + processing_method = "Azure OpenAI" + else: + result = visualizer_from_content(content, mime_type, question) + processing_method = "OpenAI" + + logger.info(f"Successfully processed image '{filename}' using {processing_method}") + return result + + except Exception as e: + logger.error(f"Error processing image '{filename}': {e!s}") + return f"Error processing image '{filename}': {e!s}" diff --git a/tests/test_visual_qa.py b/tests/test_visual_qa.py index 8fbfd71..056e79c 100644 --- a/tests/test_visual_qa.py +++ b/tests/test_visual_qa.py @@ -1,22 +1,22 @@ import base64 import os import tempfile -from unittest.mock import Mock, mock_open, patch +from unittest.mock import Mock, patch import pytest from PIL import Image from mxtoai.scripts.visual_qa import ( VisualQATool, - azure_visualizer, - encode_image, - resize_image, - visualizer, + azure_visualizer_from_content, + encode_image_from_content, + resize_image_from_content, + visualizer_from_content, ) class TestEncodeImage: - """Test the encode_image function.""" + """Test the encode_image_from_content function.""" def test_encode_local_image_file(self): """Test encoding a local image file.""" @@ -28,7 +28,11 @@ def test_encode_local_image_file(self): temp_path = temp_file.name try: - result = encode_image(temp_path) + # Read the file content + with open(temp_path, "rb") as f: + content = f.read() + + result = encode_image_from_content(content) # Should return a base64 string assert isinstance(result, str) @@ -41,79 +45,60 @@ def test_encode_local_image_file(self): finally: os.unlink(temp_path) - @patch("requests.get") - def test_encode_remote_image_url(self, mock_get): - """Test encoding an image from URL.""" - # Mock HTTP response - mock_response = Mock() - mock_response.raise_for_status.return_value = None - mock_response.headers = {"content-type": "image/jpeg"} - mock_response.iter_content.return_value = [b"fake_jpeg_data"] - mock_get.return_value = mock_response + def test_encode_image_from_bytes(self): + """Test encoding image from bytes.""" + # Create a simple test image in memory + img = Image.new("RGB", (50, 50), color="green") + import io - # Mock file operations - provide binary data for base64 encoding - with patch("builtins.open", mock_open(read_data=b"fake_jpeg_data")) as mock_file, patch("os.path.abspath") as mock_abspath: - mock_abspath.return_value = "/downloads/test.jpg" + img_bytes = io.BytesIO() + img.save(img_bytes, format="JPEG") + content = img_bytes.getvalue() - result = encode_image("http://example.com/image.jpg") + result = encode_image_from_content(content) - assert isinstance(result, str) - mock_get.assert_called_once() - mock_file.assert_called() - - @patch("requests.get") - def test_encode_remote_image_no_extension(self, mock_get): - """Test encoding remote image when content-type doesn't map to extension.""" - mock_response = Mock() - mock_response.raise_for_status.return_value = None - mock_response.headers = {"content-type": "application/octet-stream"} - mock_response.iter_content.return_value = [b"data"] - mock_get.return_value = mock_response + # Should return a base64 string + assert isinstance(result, str) + assert len(result) > 0 - with patch("builtins.open", mock_open(read_data=b"data")): - with patch("os.path.abspath", return_value="/downloads/test.download"): - # Should handle unknown content types - result = encode_image("http://example.com/unknown") - assert isinstance(result, str) + # Should be valid base64 + decoded = base64.b64decode(result) + assert len(decoded) > 0 def test_encode_nonexistent_file(self): - """Test encoding a file that doesn't exist.""" - with pytest.raises(FileNotFoundError): - encode_image("/nonexistent/file.jpg") + """Test encoding with empty content.""" + # encode_image_from_content just does base64 encoding, so empty bytes will return empty string + result = encode_image_from_content(b"") + assert result == "" class TestResizeImage: - """Test the resize_image function.""" + """Test the resize_image_from_content function.""" def test_resize_image_success(self): """Test successful image resizing.""" - # Create a temporary image - with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file: - img = Image.new("RGB", (200, 200), color="blue") - img.save(temp_file.name, "JPEG") - temp_path = temp_file.name + # Create a test image in memory + img = Image.new("RGB", (200, 200), color="blue") + import io - resized_path = None - try: - resized_path = resize_image(temp_path) + img_bytes = io.BytesIO() + img.save(img_bytes, format="JPEG") + content = img_bytes.getvalue() - # Check that resized file was created - assert os.path.exists(resized_path) - assert "resized_" in resized_path + resized_content = resize_image_from_content(content) - # Check dimensions are halved - resized_img = Image.open(resized_path) - assert resized_img.size == (100, 100) + # Check that we got bytes back + assert isinstance(resized_content, bytes) + assert len(resized_content) > 0 - finally: - os.unlink(temp_path) - if resized_path and os.path.exists(resized_path): - os.unlink(resized_path) + # Check dimensions are halved + resized_img = Image.open(io.BytesIO(resized_content)) + assert resized_img.size == (100, 100) def test_resize_nonexistent_file(self): - """Test resizing a file that doesn't exist.""" - with pytest.raises(FileNotFoundError): - resize_image("/nonexistent/file.jpg") + """Test resizing with empty content.""" + with pytest.raises(Exception): # PIL will raise UnidentifiedImageError + resize_image_from_content(b"") class TestProcessImagesAndText: @@ -133,29 +118,30 @@ def test_initialization(self): assert tool.name == "visualizer" assert "answer questions about attached images" in tool.description - assert "image_path" in tool.inputs + assert "content" in tool.inputs + assert "mime_type" in tool.inputs assert "question" in tool.inputs assert tool.output_type == "string" assert hasattr(tool, "client") - @patch("mxtoai.scripts.visual_qa.process_images_and_text") + @patch("mxtoai.scripts.visual_qa.process_images_and_text_from_content") def test_forward_with_question(self, mock_process): """Test forward method with a specific question.""" mock_process.return_value = {"generated_text": "This is a test image"} tool = VisualQATool() - result = tool.forward("test_image.jpg", "What do you see?") + result = tool.forward(b"fake_image_content", "image/jpeg", "What do you see?") assert result == {"generated_text": "This is a test image"} - mock_process.assert_called_once_with("test_image.jpg", "What do you see?", tool.client) + mock_process.assert_called_once_with(b"fake_image_content", "What do you see?", tool.client) - @patch("mxtoai.scripts.visual_qa.process_images_and_text") + @patch("mxtoai.scripts.visual_qa.process_images_and_text_from_content") def test_forward_without_question(self, mock_process): """Test forward method without a specific question (auto-caption).""" mock_process.return_value = {"generated_text": "Auto-generated caption"} tool = VisualQATool() - result = tool.forward("test_image.jpg") + result = tool.forward(b"fake_image_content", "image/jpeg") # Should add explanatory note assert "You did not provide a particular question" in result @@ -166,22 +152,19 @@ def test_forward_without_question(self, mock_process): call_args = mock_process.call_args assert "Please write a detailed caption" in call_args[0][1] - @patch("mxtoai.scripts.visual_qa.process_images_and_text") - @patch("mxtoai.scripts.visual_qa.resize_image") + @patch("mxtoai.scripts.visual_qa.process_images_and_text_from_content") + @patch("mxtoai.scripts.visual_qa.resize_image_from_content") def test_forward_payload_too_large_retry(self, mock_resize, mock_process): """Test handling of 'Payload Too Large' error with image resizing.""" # First call fails with payload error, second succeeds - mock_process.side_effect = [ - Exception("Payload Too Large"), - {"generated_text": "Resized image result"} - ] - mock_resize.return_value = "resized_test_image.jpg" + mock_process.side_effect = [Exception("Payload Too Large"), {"generated_text": "Resized image result"}] + mock_resize.return_value = b"resized_image_content" tool = VisualQATool() - result = tool.forward("test_image.jpg", "What's this?") + result = tool.forward(b"fake_image_content", "image/jpeg", "What's this?") assert result == {"generated_text": "Resized image result"} - mock_resize.assert_called_once_with("test_image.jpg") + mock_resize.assert_called_once_with(b"fake_image_content") assert mock_process.call_count == 2 @@ -189,30 +172,24 @@ class TestVisualizerFunction: """Test the visualizer function.""" def test_invalid_image_path_type(self): - """Test visualizer with invalid image_path type.""" - with pytest.raises(Exception, match="You should provide at least `image_path` string argument"): - visualizer(123) # Non-string path + """Test visualizer with invalid content type.""" + with pytest.raises(TypeError): + visualizer_from_content(123, "image/jpeg") # Non-bytes content @patch("requests.post") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_visualizer_success(self, mock_encode, mock_post): """Test successful image analysis with visualizer function.""" mock_encode.return_value = "base64_encoded_image" mock_response = Mock() - mock_response.json.return_value = { - "choices": [{ - "message": { - "content": "This is an image of a cat" - } - }] - } + mock_response.json.return_value = {"choices": [{"message": {"content": "This is an image of a cat"}}]} mock_post.return_value = mock_response - result = visualizer("test_image.jpg", "What animal is this?") + result = visualizer_from_content(b"test_image_content", "image/jpeg", "What animal is this?") assert result == "This is an image of a cat" - mock_encode.assert_called_once_with("test_image.jpg") + mock_encode.assert_called_once_with(b"test_image_content") # Check API call mock_post.assert_called_once() @@ -231,22 +208,16 @@ def test_visualizer_success(self, mock_encode, mock_post): assert content[1]["type"] == "image_url" @patch("requests.post") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_visualizer_without_question(self, mock_encode, mock_post): """Test visualizer function without specific question.""" mock_encode.return_value = "base64_data" mock_response = Mock() - mock_response.json.return_value = { - "choices": [{ - "message": { - "content": "Auto-generated description" - } - }] - } + mock_response.json.return_value = {"choices": [{"message": {"content": "Auto-generated description"}}]} mock_post.return_value = mock_response - result = visualizer("image.jpg") + result = visualizer_from_content(b"image_content", "image/jpeg") # Should add explanatory note assert "You did not provide a particular question" in result @@ -259,7 +230,7 @@ def test_visualizer_without_question(self, mock_encode, mock_post): assert "Please write a detailed caption" in content @patch("requests.post") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_visualizer_api_error(self, mock_encode, mock_post): """Test visualizer function handling API errors.""" mock_encode.return_value = "base64_data" @@ -269,20 +240,23 @@ def test_visualizer_api_error(self, mock_encode, mock_post): mock_post.return_value = mock_response with pytest.raises(Exception, match="Response format unexpected"): - visualizer("image.jpg", "What's this?") + visualizer_from_content(b"image_content", "image/jpeg", "What's this?") class TestAzureVisualizer: """Test the azure_visualizer function.""" @patch("mxtoai.scripts.visual_qa.completion") - @patch("mxtoai.scripts.visual_qa.encode_image") - @patch.dict(os.environ, { - "MODEL_NAME": "gpt-4o", - "MODEL_API_KEY": "test_key", - "MODEL_ENDPOINT": "https://test.openai.azure.com", - "MODEL_API_VERSION": "2024-02-01" - }) + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") + @patch.dict( + os.environ, + { + "MODEL_NAME": "gpt-4o", + "MODEL_API_KEY": "test_key", + "MODEL_ENDPOINT": "https://test.openai.azure.com", + "MODEL_API_VERSION": "2024-02-01", + }, + ) def test_azure_visualizer_success(self, mock_encode, mock_completion): """Test successful Azure OpenAI image analysis.""" mock_encode.return_value = "base64_image_data" @@ -293,10 +267,10 @@ def test_azure_visualizer_success(self, mock_encode, mock_completion): mock_response.choices[0].message.content = "Azure analysis result" mock_completion.return_value = mock_response - result = azure_visualizer("test_image.jpg", "Analyze this image") + result = azure_visualizer_from_content(b"test_image_content", "image/jpeg", "Analyze this image") assert result == "Azure analysis result" - mock_encode.assert_called_once_with("test_image.jpg") + mock_encode.assert_called_once_with(b"test_image_content") # Check litellm completion call mock_completion.assert_called_once() @@ -320,7 +294,7 @@ def test_azure_visualizer_success(self, mock_encode, mock_completion): assert content[1]["type"] == "image_url" @patch("mxtoai.scripts.visual_qa.completion") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_azure_visualizer_without_question(self, mock_encode, mock_completion): """Test Azure visualizer without specific question.""" mock_encode.return_value = "base64_data" @@ -331,14 +305,14 @@ def test_azure_visualizer_without_question(self, mock_encode, mock_completion): mock_response.choices[0].message.content = "Default caption" mock_completion.return_value = mock_response - result = azure_visualizer("image.jpg") + result = azure_visualizer_from_content(b"image_content", "image/jpeg") assert "You did not provide a particular question" in result assert "detailed caption for the image" in result @patch("mxtoai.scripts.visual_qa.completion") - @patch("mxtoai.scripts.visual_qa.encode_image") - @patch("mxtoai.scripts.visual_qa.resize_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") + @patch("mxtoai.scripts.visual_qa.resize_image_from_content") def test_azure_visualizer_image_too_large(self, mock_resize, mock_encode, mock_completion): """Test Azure visualizer handling large image errors.""" mock_encode.return_value = "base64_data" @@ -347,17 +321,17 @@ def test_azure_visualizer_image_too_large(self, mock_resize, mock_encode, mock_c # First call fails with size error, second succeeds after resize mock_completion.side_effect = [ Exception("image too large"), - Mock(choices=[Mock(message=Mock(content="Resized result"))]) + Mock(choices=[Mock(message=Mock(content="Resized result"))]), ] - result = azure_visualizer("large_image.jpg", "What's this?") + result = azure_visualizer_from_content(b"large_image_content", "image/jpeg", "What's this?") # The function successfully retries with resized image and returns the result assert result == "Resized result" - mock_resize.assert_called_once_with("large_image.jpg") + mock_resize.assert_called_once_with(b"large_image_content") @patch("mxtoai.scripts.visual_qa.completion") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_azure_visualizer_empty_response(self, mock_encode, mock_completion): """Test Azure visualizer handling empty response.""" mock_encode.return_value = "base64_data" @@ -368,23 +342,20 @@ def test_azure_visualizer_empty_response(self, mock_encode, mock_completion): mock_response.choices[0].message.content = "" mock_completion.return_value = mock_response - result = azure_visualizer("image.jpg", "What's this?") - - # The function catches the exception and returns an error message - assert "Error processing image: Empty response from Azure OpenAI" in result + with pytest.raises(Exception, match="Failed to process image"): + azure_visualizer_from_content(b"image_content", "image/jpeg", "What's this?") @patch("mxtoai.scripts.visual_qa.completion") - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_azure_visualizer_api_error(self, mock_encode, mock_completion): """Test Azure visualizer handling API errors.""" mock_encode.return_value = "base64_data" mock_completion.side_effect = Exception("API connection failed") - result = azure_visualizer("image.jpg", "What's this?") + with pytest.raises(Exception, match="Failed to process image"): + azure_visualizer_from_content(b"image_content", "image/jpeg", "What's this?") - assert "Error processing image: API connection failed" in result - - @patch("mxtoai.scripts.visual_qa.encode_image") + @patch("mxtoai.scripts.visual_qa.encode_image_from_content") def test_azure_visualizer_mime_type_detection(self, mock_encode): """Test MIME type detection for different image formats.""" mock_encode.return_value = "base64_data" @@ -397,15 +368,10 @@ def test_azure_visualizer_mime_type_detection(self, mock_encode): mock_completion.return_value = mock_response # Test different image extensions - test_files = [ - "image.jpg", - "image.png", - "image.gif", - "image.unknown" - ] + test_files = ["image.jpg", "image.png", "image.gif", "image.unknown"] for filename in test_files: - result = azure_visualizer(filename, "Analyze") + result = azure_visualizer_from_content(b"image_content", "image/jpeg", "Analyze") assert result == "Analysis result" # Check MIME type handling @@ -434,30 +400,32 @@ def test_visual_qa_tool_integration(self): temp_path = temp_file.name try: - with patch("mxtoai.scripts.visual_qa.process_images_and_text") as mock_process: + with patch("mxtoai.scripts.visual_qa.process_images_and_text_from_content") as mock_process: mock_process.return_value = {"generated_text": "Integration test result"} - result = tool.forward(temp_path, "What color is this image?") + # Read the file content to pass as bytes + with open(temp_path, "rb") as f: + content = f.read() + + result = tool.forward(content, "image/jpeg", "What color is this image?") # Should return the processed result assert result == {"generated_text": "Integration test result"} - mock_process.assert_called_once_with( - temp_path, - "What color is this image?", - tool.client - ) + mock_process.assert_called_once_with(content, "What color is this image?", tool.client) finally: os.unlink(temp_path) def test_error_handling_chain(self): """Test error handling across different visual processing functions.""" - # Test that errors propagate properly through the processing chain - with pytest.raises(FileNotFoundError): - encode_image("nonexistent_file.jpg") - - with pytest.raises(FileNotFoundError): - resize_image("nonexistent_file.jpg") + # encode_image_from_content just does base64 encoding, won't raise FileNotFoundError + result = encode_image_from_content(b"") + assert result == "" + # resize_image_from_content will raise an image processing error with pytest.raises(Exception): - visualizer(None) # Invalid input type + resize_image_from_content(b"") + + # visualizer_from_content requires proper parameters + with pytest.raises(TypeError): + visualizer_from_content(None, "image/jpeg") # Invalid input type diff --git a/tests/test_visual_qa_tool.py b/tests/test_visual_qa_tool.py new file mode 100644 index 0000000..139e58c --- /dev/null +++ b/tests/test_visual_qa_tool.py @@ -0,0 +1,193 @@ +import io +from unittest.mock import Mock, patch + +from PIL import Image + +from mxtoai.tools.visual_qa_tool import VisualQATool + + +class TestVisualQATool: + """Test the VisualQATool class integration with attachment service.""" + + def test_initialization(self): + """Test VisualQATool initialization.""" + tool = VisualQATool() + + assert tool.name == "visual_qa" + assert "Answer questions about image attachments" in tool.description + assert "filename" in tool.inputs + assert "question" in tool.inputs + assert "use_azure" in tool.inputs + assert tool.output_type == "string" + + def test_no_attachment_service(self): + """Test error when no attachment service is available.""" + tool = VisualQATool() # No context provided + + result = tool.forward("test.jpg", "What is this?") + + assert "Error: No attachment service available" in result + + def test_attachment_not_found(self): + """Test error when attachment is not found.""" + # Mock context with attachment service + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = False + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("nonexistent.jpg", "What is this?") + + assert "Error: Attachment 'nonexistent.jpg' not found" in result + + @patch("mxtoai.tools.visual_qa_tool.azure_visualizer_from_content") + def test_successful_azure_processing(self, mock_azure_visualizer): + """Test successful image processing with Azure OpenAI.""" + mock_azure_visualizer.return_value = "This is a test image showing a red square." + + # Create a simple test image in memory + img = Image.new("RGB", (100, 100), color="red") + img_bytes = io.BytesIO() + img.save(img_bytes, format="JPEG") + content = img_bytes.getvalue() + + # Mock context with attachment service + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = content + mock_attachment_service.get_metadata.return_value = {"type": "image/jpeg", "size": len(content)} + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("test.jpg", "What color is this?", use_azure=True) + + assert "This is a test image showing a red square." in result + mock_azure_visualizer.assert_called_once_with(content, "image/jpeg", "What color is this?") + + @patch("mxtoai.tools.visual_qa_tool.visualizer_from_content") + def test_successful_openai_processing(self, mock_visualizer): + """Test successful image processing with standard OpenAI.""" + mock_visualizer.return_value = "This is a test image showing a blue circle." + + # Create a simple test image in memory + img = Image.new("RGB", (100, 100), color="blue") + img_bytes = io.BytesIO() + img.save(img_bytes, format="PNG") + content = img_bytes.getvalue() + + # Mock context with attachment service + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = content + mock_attachment_service.get_metadata.return_value = {"type": "image/png", "size": len(content)} + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("test.png", "What shape is this?", use_azure=False) + + assert "This is a test image showing a blue circle." in result + mock_visualizer.assert_called_once_with(content, "image/png", "What shape is this?") + + def test_non_image_file(self): + """Test error when trying to process non-image file.""" + # Mock context with attachment service returning a text file + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = b"This is text content" + mock_attachment_service.get_metadata.return_value = {"type": "text/plain", "size": 20} + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("document.txt", "What is this?") + + assert "is not an image file" in result + + def test_mime_type_detection(self): + """Test MIME type detection for files without proper metadata.""" + # Create a simple test image + img = Image.new("RGB", (50, 50), color="green") + img_bytes = io.BytesIO() + img.save(img_bytes, format="JPEG") + content = img_bytes.getvalue() + + # Mock context with attachment service - no type in metadata + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = content + mock_attachment_service.get_metadata.return_value = {"size": len(content)} # No type field + mock_context.attachment_service = mock_attachment_service + + with patch("mxtoai.tools.visual_qa_tool.azure_visualizer_from_content") as mock_azure: + mock_azure.return_value = "Green square detected" + + tool = VisualQATool(context=mock_context) + result = tool.forward("image.jpg", "What color?") + + assert "Green square detected" in result + # Should detect JPEG from filename + mock_azure.assert_called_once_with(content, "image/jpeg", "What color?") + + @patch("mxtoai.tools.visual_qa_tool.azure_visualizer_from_content") + def test_processing_exception_handling(self, mock_azure_visualizer): + """Test error handling when image processing fails.""" + mock_azure_visualizer.side_effect = Exception("API error") + + # Create test image + img = Image.new("RGB", (100, 100), color="red") + img_bytes = io.BytesIO() + img.save(img_bytes, format="JPEG") + content = img_bytes.getvalue() + + # Mock context + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = content + mock_attachment_service.get_metadata.return_value = {"type": "image/jpeg", "size": len(content)} + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("test.jpg", "What is this?") + + assert "Error processing image 'test.jpg': API error" in result + + def test_content_retrieval_failure(self): + """Test error when content cannot be retrieved.""" + # Mock context with attachment service that can't retrieve content + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = None # Content retrieval fails + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("test.jpg", "What is this?") + + assert "Error: Could not retrieve content for 'test.jpg'" in result + + def test_metadata_retrieval_failure(self): + """Test error when metadata cannot be retrieved.""" + # Mock context with attachment service that can't retrieve metadata + mock_context = Mock() + mock_attachment_service = Mock() + mock_attachment_service.has_attachment.return_value = True + mock_attachment_service.get_content.return_value = b"image_content" + mock_attachment_service.get_metadata.return_value = None # Metadata retrieval fails + mock_context.attachment_service = mock_attachment_service + + tool = VisualQATool(context=mock_context) + + result = tool.forward("test.jpg", "What is this?") + + assert "Error: Could not retrieve metadata for 'test.jpg'" in result