Update main.py

main.py

@@ -132,19 +132,24 @@ def update_db(user_id, conversation_id, message, response):
 # Utility Functions
 # ============================================================================
 
-def extract_data_from_tag(input_string, tag):
+def extract_data_from_tag(input_string, tag, invert=False):
     # Create the regex pattern
     pattern = f'<{tag}.*?>(.*?)</{tag}>'
     
     # Find all matches
     matches = re.findall(pattern, input_string, re.DOTALL)
     
-    # If matches are found, return them joined by newlines
-    if matches:
-        out = '\n'.join(match.strip() for match in matches)
-        return out
+    if invert:
+        # If inverted, replace all matches with empty string
+        out = re.sub(pattern, '', input_string, flags=re.DOTALL)
+        return out.strip()
     else:
-        return input_string
+        # If matches are found, return them joined by newlines
+        if matches:
+            out = '\n'.join(match.strip() for match in matches)
+            return out
+        else:
+            return input_string
 
 def calculate_tokens(msgs):
     return sum(len(encoding.encode(str(m))) for m in msgs)
@@ -409,6 +414,9 @@ async def presentation_chat(query: PresentationChatModel, background_tasks: Back
 
     # Collect the entire response
     full_response = ""
+    html_content = ""
+    marp_content_with_images = ""
+    
     async for chunk in response_stream.body_iterator:
         full_response += chunk
 
@@ -416,16 +424,18 @@ async def presentation_chat(query: PresentationChatModel, background_tasks: Back
     # Extract the Marp presentation content
     marp_content = extract_data_from_tag(full_response, "marp_presentation")
 
-    # Replace image keywords
-    marp_content_with_images = replace_image_keywords(marp_content)
+    if marp_content:
+        # Replace image keywords
+        marp_content_with_images = replace_image_keywords(marp_content)
+    
+        # Convert Markdown to HTML
+        html_content = convert_markdown_marp(marp_content_with_images, 'html')
 
-    # Convert Markdown to HTML
-    html_content = convert_markdown_marp(marp_content_with_images, 'html')
 
-    if html_content:
-        return JSONResponse({
-            "markdown": marp_content_with_images,
-            "html": html_content.decode()
+    return JSONResponse({
+            "response": extract_data_from_tag(full_response, "marp_presentation",invert=True),
+            "markdown_presentation": marp_content_with_images,
+            "html_presentation": html_content.decode() if isinstance(html_content, bytes) else html_content
         })
     else:
         raise HTTPException(status_code=500, detail="Failed to create presentation.")