Update app.py

app.py

@@ -173,7 +173,7 @@ def chunk_markdown(
     chunk_overlap: int = 200,
     strip_headers: bool = True
 ) -> List[Document]:
-    """Chunks markdown text, preserving headers in metadata."""
+    """Chunks markdown text, preserving headers in metadata and extracting base64 images."""
     if not markdown_text_with_images or not markdown_text_with_images.strip():
         logger.warning("chunk_markdown received empty input.")
         return []
@@ -188,12 +188,15 @@ def chunk_markdown(
     header_chunks = markdown_splitter.split_text(markdown_text_with_images)
 
     if not header_chunks:
+        logger.warning("No chunks created from markdown splitting.")
         return []
 
     final_chunks = []
     if chunk_size > 0:
         text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len,
+            chunk_size=chunk_size, 
+            chunk_overlap=chunk_overlap, 
+            length_function=len,
             separators=["\n\n", "\n", "(?<=\. )", "(?<=\? )", "(?<=! )", ", ", "; ", " ", ""],
             add_start_index=True
         )
@@ -206,17 +209,23 @@ def chunk_markdown(
     else:
         final_chunks = [chunk for chunk in header_chunks if chunk.page_content]
 
+    # Extract base64 images and add to metadata
     for chunk in final_chunks:
+        if not hasattr(chunk, 'metadata'):
+            chunk.metadata = {}
+        
+        # Improved regex to capture full base64 data URI
         images_in_chunk = re.findall(
-            r"!\[.*?\]\((data:image/[a-zA-Z+]+;base64,[A-Za-z0-9+/=]+)\)",
+            r"!\[.*?\]\((data:image/[a-zA-Z]+;base64,[A-Za-z0-9+/]+={0,2})\)",
             chunk.page_content
         )
-        if not hasattr(chunk, 'metadata'):
-            chunk.metadata = {}
-        chunk.metadata["images_base64"] = images_in_chunk
+        chunk.metadata["images_base64"] = images_in_chunk if images_in_chunk else []
+        logger.debug(f"Chunk metadata updated with {len(images_in_chunk)} base64 images")
 
+    logger.info(f"Created {len(final_chunks)} chunks with base64 metadata")
     return final_chunks
 
+
 def get_hf_token(explicit_token: str = None) -> str:
     """Retrieve Hugging Face token with fallback mechanisms."""
     global hf_token_global