Create utils/response_formatter.py

utils/response_formatter.py

@@ -0,0 +1,146 @@
+# utils/response_formatter.py
+
+import re
+from typing import Dict, Optional
+import streamlit as st
+from datetime import datetime
+
+class ResponseFormatter:
+    def __init__(self):
+        self.section_keywords = [
+            "Summary",
+            "Overview",
+            "Background",
+            "Solution",
+            "Approach",
+            "Benefits",
+            "Experience",
+            "Methodology",
+            "Implementation",
+            "Timeline",
+            "Pricing",
+            "Why Us",
+            "Next Steps"
+        ]
+    
+    def format_response(self, content: str, metadata: Optional[Dict] = None) -> str:
+        """
+        Format the AI response into a clean, structured format.
+        
+        Args:
+            content (str): Raw response content
+            metadata (Optional[Dict]): Additional metadata like sources
+            
+        Returns:
+            str: Formatted response
+        """
+        # Clean up the content
+        formatted = self._clean_content(content)
+        
+        # Add structural formatting
+        formatted = self._add_structure(formatted)
+        
+        # Format lists and bullet points
+        formatted = self._format_lists(formatted)
+        
+        # Add source references if available
+        if metadata and 'sources' in metadata:
+            formatted = self._add_sources(formatted, metadata['sources'])
+        
+        return formatted
+    
+    def _clean_content(self, content: str) -> str:
+        """Clean and normalize the content."""
+        # Remove multiple newlines
+        cleaned = re.sub(r'\n{3,}', '\n\n', content)
+        
+        # Ensure consistent heading formatting
+        for keyword in self.section_keywords:
+            pattern = rf'(?i)({keyword}:?)\s*\n'
+            cleaned = re.sub(pattern, f'### {keyword}\n\n', cleaned)
+        
+        return cleaned
+    
+    def _add_structure(self, content: str) -> str:
+        """Add structural elements to the content."""
+        # Add section breaks
+        sections = re.split(r'(?m)^###\s+', content)
+        
+        if len(sections) == 1:  # No sections found
+            # Add a default section if none exists
+            content = "### Response\n\n" + content
+        
+        # Add copy button for each section
+        structured = ""
+        for section in sections:
+            if section.strip():
+                section_title = section.split('\n')[0].strip()
+                section_content = '\n'.join(section.split('\n')[1:]).strip()
+                
+                structured += f"### {section_title}\n\n"
+                structured += f"{section_content}\n\n"
+                
+        return structured
+    
+    def _format_lists(self, content: str) -> str:
+        """Format lists and bullet points consistently."""
+        # Format bullet points
+        content = re.sub(r'(?m)^[\-\*]\s+', '• ', content)
+        
+        # Format numbered lists
+        content = re.sub(r'(?m)^\d+\.\s+', lambda m: f"{int(m.group().split('.')[0])}. ", content)
+        
+        return content
+    
+    def _add_sources(self, content: str, sources: list) -> str:
+        """Add source references to the content."""
+        if sources:
+            content += "\n\n---\n\n### Sources\n\n"
+            for source in sources:
+                content += f"• {source}\n"
+        
+        return content
+
+def display_formatted_response(response_content: str, metadata: Optional[Dict] = None):
+    """
+    Display a formatted response in the Streamlit interface.
+    
+    Args:
+        response_content (str): Raw response content
+        metadata (Optional[Dict]): Additional metadata like sources
+    """
+    formatter = ResponseFormatter()
+    formatted_content = formatter.format_response(response_content, metadata)
+    
+    # Display formatted content
+    st.markdown(formatted_content)
+    
+    # Add copy buttons for sections
+    sections = re.split(r'(?m)^###\s+', formatted_content)
+    for section in sections:
+        if section.strip():
+            section_title = section.split('\n')[0].strip()
+            section_content = '\n'.join(section.split('\n')[1:]).strip()
+            
+            with st.expander(f"Copy {section_title}", expanded=False):
+                st.text_area(
+                    "Copy this content:",
+                    value=section_content,
+                    height=200,
+                    key=f"copy_{section_title}_{datetime.now().timestamp()}"
+                )
+                if st.button(f"Copy {section_title}", key=f"btn_{section_title}_{datetime.now().timestamp()}"):
+                    st.write("Content copied to clipboard!")
+
+# Example usage in your chat interface:
+"""
+def display_chat_message(message, is_user=False):
+    if is_user:
+        st.chat_message("user").write(message)
+    else:
+        with st.chat_message("assistant"):
+            display_formatted_response(
+                message.content,
+                metadata=getattr(message, 'metadata', None)
+            )
+"""