Upd code and diagram viewer + pdf

helpers/diagram.py

@@ -27,7 +27,7 @@ async def generate_mermaid_diagram(
     nvidia_rotator,
     render_error: str = "",
     retry: int = 0,
-    max_retries: int = 2
+    max_retries: int = 5
 ) -> str:
     from utils.api.router import generate_answer_with_model
 
@@ -40,24 +40,36 @@ async def generate_mermaid_diagram(
             overview.append(f"{section_id} {title}: {syn[:180]}...")
     context_overview = "\n".join(overview)
 
+    # Enhanced system prompt with better error handling guidance
     sys_prompt = (
-        "You are an expert technical illustrator. Create a single concise Mermaid diagram that best conveys the core structure\n"
+        "You are an expert technical illustrator and Mermaid syntax specialist. Create a single concise Mermaid diagram that best conveys the core structure\n"
         "(e.g., flowchart, sequence, class, state, or ER) based on the provided CONTEXT.\n"
         "Rules:\n"
         "- Return Mermaid code only (no backticks, no explanations).\n"
         "- Prefer flowchart or sequence if uncertain.\n"
         "- Keep node labels short but meaningful.\n"
-        "- Ensure Mermaid syntax is valid.\n"
+        "- Ensure Mermaid syntax is valid and follows these guidelines:\n"
+        "  * Use proper node IDs (alphanumeric, no spaces)\n"
+        "  * Use proper arrow syntax (--> for flowcharts, ->> for sequence)\n"
+        "  * Quote labels with special characters\n"
+        "  * Use proper diagram type declarations\n"
+        "- If there was a previous error, fix the specific syntax issues mentioned.\n"
     )
 
-    feedback = (f"\n\nRENDERING ERROR TO FIX: {render_error}\n" if render_error else "")
+    # Enhanced error feedback
+    if render_error:
+        feedback = f"\n\nPREVIOUS RENDERING ERROR TO FIX:\n{render_error}\n\nPlease analyze this error and generate a corrected Mermaid diagram that addresses the specific syntax or logical issues mentioned above."
+    else:
+        feedback = ""
+
     user_prompt = (
         f"INSTRUCTIONS:\n{instructions}\n\nCONTEXT OVERVIEW:\n{context_overview}{feedback}"
     )
 
+    # Use NVIDIA_LARGE for better diagram generation
     selection = {"provider": "nvidia_large", "model": "openai/gpt-oss-120b"}
 
-    logger.info(f"[DIAGRAM] Generating Mermaid (retry={retry})")
+    logger.info(f"[DIAGRAM] Generating Mermaid (retry={retry}/{max_retries})")
     diagram = await generate_answer_with_model(selection, sys_prompt, user_prompt, gemini_rotator, nvidia_rotator)
     diagram = (diagram or "").strip()
 
@@ -67,7 +79,7 @@ async def generate_mermaid_diagram(
         if raw.lower().startswith("mermaid"):
             diagram = "\n".join(raw.splitlines()[1:])
 
-    # Naive validation: basic mermaid keywords
+    # Enhanced validation: check for common Mermaid syntax issues
     if not any(kw in diagram for kw in ("graph", "sequenceDiagram", "classDiagram", "stateDiagram", "erDiagram")):
         logger.warning("[DIAGRAM] Mermaid validation failed: missing diagram keywords")
         if retry < max_retries:
@@ -79,3 +91,176 @@ async def generate_mermaid_diagram(
     return diagram
 
 
+async def _render_mermaid_with_retry(mermaid_text: str, max_retries: int = 3) -> bytes:
+    """
+    Render mermaid code to PNG with retry logic and AI-powered error correction.
+    """
+    last_error = ""
+    
+    for attempt in range(max_retries):
+        try:
+            # Try to render the current mermaid code
+            img_bytes = _render_mermaid_png(mermaid_text)
+            
+            if img_bytes and len(img_bytes) > 0:
+                logger.info(f"[DIAGRAM] Mermaid rendered successfully on attempt {attempt + 1}")
+                return img_bytes
+            else:
+                logger.warning(f"[DIAGRAM] Mermaid render returned empty on attempt {attempt + 1}")
+                
+        except Exception as e:
+            last_error = str(e)
+            logger.warning(f"[DIAGRAM] Mermaid render attempt {attempt + 1} failed: {e}")
+        
+        # If this isn't the last attempt, try to fix the mermaid code using AI
+        if attempt < max_retries - 1:
+            try:
+                logger.info(f"[DIAGRAM] Attempting to fix Mermaid syntax using AI (attempt {attempt + 1})")
+                fixed_mermaid = await _fix_mermaid_with_ai(mermaid_text, last_error)
+                if fixed_mermaid and fixed_mermaid != mermaid_text:
+                    mermaid_text = fixed_mermaid
+                    logger.info(f"[DIAGRAM] AI provided fixed Mermaid code for retry {attempt + 2}")
+                else:
+                    logger.warning(f"[DIAGRAM] AI could not provide fixed Mermaid code")
+                    break
+            except Exception as ai_error:
+                logger.warning(f"[DIAGRAM] AI Mermaid fix failed: {ai_error}")
+                break
+    
+    logger.warning(f"[DIAGRAM] All Mermaid render attempts failed, last error: {last_error}")
+    return b""
+
+
+async def _fix_mermaid_with_ai(mermaid_text: str, error_message: str) -> str:
+    """
+    Use AI to fix Mermaid syntax errors.
+    """
+    try:
+        from utils.api.router import generate_answer_with_model
+        
+        sys_prompt = """You are a Mermaid syntax expert. Your task is to fix Mermaid diagram syntax errors.
+
+Rules:
+1. Return ONLY the corrected Mermaid code (no backticks, no explanations)
+2. Ensure proper syntax for the diagram type
+3. Fix common issues like:
+   - Invalid node IDs (use alphanumeric, no spaces)
+   - Incorrect arrow syntax
+   - Missing quotes around labels with special characters
+   - Wrong diagram type declarations
+4. Maintain the original intent and structure of the diagram"""
+
+        user_prompt = f"""Fix this Mermaid diagram that has rendering errors:
+
+ORIGINAL MERMAID CODE:
+```mermaid
+{mermaid_text}
+```
+
+ERROR MESSAGE:
+{error_message}
+
+Please provide the corrected Mermaid code that will render successfully."""
+
+        # Use NVIDIA_LARGE for better error correction
+        selection = {"provider": "nvidia_large", "model": "openai/gpt-oss-120b"}
+        response = await generate_answer_with_model(selection, sys_prompt, user_prompt, None, None)
+        
+        if response:
+            # Clean up the response
+            fixed_code = response.strip()
+            if fixed_code.startswith("```"):
+                fixed_code = fixed_code.strip('`')
+                if fixed_code.lower().startswith("mermaid"):
+                    fixed_code = "\n".join(fixed_code.splitlines()[1:])
+            
+            return fixed_code.strip()
+        
+    except Exception as e:
+        logger.warning(f"[DIAGRAM] AI Mermaid fix failed: {e}")
+    
+    return ""
+
+
+def _render_mermaid_png(mermaid_text: str) -> bytes:
+    """
+    Render mermaid code to PNG via Kroki service (no local mermaid-cli dependency).
+    Falls back to returning empty bytes on failure.
+    """
+    try:
+        import base64
+        import json
+        import urllib.request
+        import urllib.error
+        
+        # Validate and clean mermaid content
+        if not mermaid_text or not mermaid_text.strip():
+            logger.warning("[DIAGRAM] Empty mermaid content")
+            return b""
+        
+        # Clean the mermaid text - remove any potential issues
+        cleaned_text = mermaid_text.strip()
+        
+        # Basic mermaid syntax validation
+        if not cleaned_text.startswith(('graph', 'flowchart', 'sequenceDiagram', 'classDiagram', 'stateDiagram', 'erDiagram', 'journey', 'gantt', 'pie', 'gitgraph')):
+            logger.warning(f"[DIAGRAM] Invalid mermaid diagram type: {cleaned_text[:50]}...")
+            return b""
+        
+        # Kroki POST API for mermaid -> png
+        data = json.dumps({"diagram_source": cleaned_text}).encode("utf-8")
+        req = urllib.request.Request(
+            url="https://kroki.io/mermaid/png",
+            data=data,
+            headers={"Content-Type": "application/json"},
+            method="POST"
+        )
+        
+        with urllib.request.urlopen(req, timeout=15) as resp:
+            if resp.status == 200:
+                return resp.read()
+            else:
+                logger.warning(f"[DIAGRAM] Kroki returned status {resp.status}")
+                return b""
+                
+    except urllib.error.HTTPError as e:
+        if e.code == 400:
+            logger.warning(f"[DIAGRAM] Kroki mermaid syntax error (400): {e.reason}")
+        else:
+            logger.warning(f"[DIAGRAM] Kroki HTTP error {e.code}: {e.reason}")
+    except urllib.error.URLError as e:
+        logger.warning(f"[DIAGRAM] Kroki connection error: {e.reason}")
+    except Exception as e:
+        logger.warning(f"[DIAGRAM] Kroki mermaid render error: {e}")
+    
+    return b""
+
+
+async def fix_mermaid_syntax_for_ui(mermaid_text: str, error_message: str = "") -> str:
+    """
+    Fix Mermaid syntax for UI rendering using AI.
+    Returns the corrected Mermaid code that can be used in the browser.
+    """
+    try:
+        # If no error message provided, try to validate the mermaid syntax first
+        if not error_message:
+            # Basic validation - check for common issues
+            if not mermaid_text.strip():
+                error_message = "Empty Mermaid diagram"
+            elif not any(kw in mermaid_text for kw in ("graph", "sequenceDiagram", "classDiagram", "stateDiagram", "erDiagram")):
+                error_message = "Missing valid Mermaid diagram type declaration"
+        
+        # Use AI to fix the mermaid code
+        fixed_code = await _fix_mermaid_with_ai(mermaid_text, error_message)
+        
+        if fixed_code and fixed_code != mermaid_text:
+            logger.info(f"[DIAGRAM] AI provided fixed Mermaid code for UI")
+            return fixed_code
+        else:
+            logger.warning(f"[DIAGRAM] AI could not fix Mermaid code for UI")
+            return mermaid_text  # Return original if AI couldn't fix it
+            
+    except Exception as e:
+        logger.warning(f"[DIAGRAM] Mermaid UI fix failed: {e}")
+        return mermaid_text  # Return original on error
+
+

routes/reports.py

@@ -1296,3 +1296,32 @@ Return the renumbered headings in the format: "level: new_number: heading_text"
         return report
 
 
+@app.post("/mermaid/fix")
+async def fix_mermaid_syntax(
+    mermaid_code: str = Form(...),
+    error_message: str = Form("")
+):
+    """
+    Fix Mermaid diagram syntax using AI for UI rendering.
+    """
+    try:
+        from helpers.diagram import fix_mermaid_syntax_for_ui
+        
+        logger.info(f"[MERMAID] Fixing Mermaid syntax for UI")
+        fixed_code = await fix_mermaid_syntax_for_ui(mermaid_code, error_message)
+        
+        return {
+            "success": True,
+            "fixed_code": fixed_code,
+            "was_fixed": fixed_code != mermaid_code
+        }
+        
+    except Exception as e:
+        logger.error(f"[MERMAID] Failed to fix Mermaid syntax: {e}")
+        return {
+            "success": False,
+            "error": str(e),
+            "fixed_code": mermaid_code
+        }
+
+

static/script.js

@@ -821,33 +821,91 @@
     const isV10 = !!(window.mermaid && window.mermaid.render && typeof window.mermaid.render === 'function');
     for (let idx = 0; idx < mermaidBlocks.length; idx++) {
       const codeBlock = mermaidBlocks[idx];
-      const graph = codeBlock.textContent || '';
+      let graph = codeBlock.textContent || '';
       const wrapper = document.createElement('div');
       const id = `mermaid-${Date.now()}-${idx}`;
       wrapper.className = 'mermaid';
       wrapper.id = id;
       const replaceTarget = codeBlock.parentElement && codeBlock.parentElement.tagName.toLowerCase() === 'pre' ? codeBlock.parentElement : codeBlock;
       replaceTarget.replaceWith(wrapper);
-      try {
-        if (isV10) {
-          // Pass wrapper as container to avoid document.createElementNS undefined errors
-          const out = await window.mermaid.render(id + '-svg', graph, wrapper);
-          if (out && out.svg) {
-            wrapper.innerHTML = out.svg;
-            if (out.bindFunctions) { out.bindFunctions(wrapper); }
+      
+      // Try to render with retry logic
+      let renderSuccess = false;
+      let attempt = 0;
+      const maxAttempts = 3;
+      
+      while (!renderSuccess && attempt < maxAttempts) {
+        try {
+          if (isV10) {
+            // Pass wrapper as container to avoid document.createElementNS undefined errors
+            const out = await window.mermaid.render(id + '-svg', graph, wrapper);
+            if (out && out.svg) {
+              wrapper.innerHTML = out.svg;
+              if (out.bindFunctions) { out.bindFunctions(wrapper); }
+              renderSuccess = true;
+            }
+          } else if (window.mermaid && window.mermaid.init) {
+            // Legacy fallback
+            wrapper.textContent = graph;
+            window.mermaid.init(undefined, wrapper);
+            renderSuccess = true;
+          }
+        } catch (e) {
+          console.warn(`Mermaid render failed (attempt ${attempt + 1}):`, e);
+          
+          // If this isn't the last attempt, try to fix the mermaid code using AI
+          if (attempt < maxAttempts - 1) {
+            try {
+              console.log('Attempting to fix Mermaid syntax using AI...');
+              const fixedCode = await fixMermaidWithAI(graph, e.message || e.toString());
+              if (fixedCode && fixedCode !== graph) {
+                graph = fixedCode;
+                console.log('AI provided fixed Mermaid code, retrying...');
+              } else {
+                console.warn('AI could not provide fixed Mermaid code');
+                break;
+              }
+            } catch (aiError) {
+              console.warn('AI Mermaid fix failed:', aiError);
+              break;
+            }
           }
-        } else if (window.mermaid && window.mermaid.init) {
-          // Legacy fallback
-          wrapper.textContent = graph;
-          window.mermaid.init(undefined, wrapper);
         }
-      } catch (e) {
-        console.warn('Mermaid render failed:', e);
+        attempt++;
+      }
+      
+      // If all attempts failed, show the original code
+      if (!renderSuccess) {
+        console.warn('All Mermaid render attempts failed, showing original code');
         wrapper.textContent = graph;
       }
     }
   }
 
+  async function fixMermaidWithAI(mermaidCode, errorMessage) {
+    try {
+      const formData = new FormData();
+      formData.append('mermaid_code', mermaidCode);
+      formData.append('error_message', errorMessage);
+      
+      const response = await fetch('/mermaid/fix', {
+        method: 'POST',
+        body: formData
+      });
+      
+      if (response.ok) {
+        const result = await response.json();
+        if (result.success && result.was_fixed) {
+          console.log('Mermaid syntax fixed by AI');
+          return result.fixed_code;
+        }
+      }
+    } catch (error) {
+      console.warn('Failed to fix Mermaid with AI:', error);
+    }
+    return null;
+  }
+
   // Expose markdown-aware appenders for use after refresh (projects.js)
   window.appendMessage = appendMessage;
   window.appendSources = appendSources;
@@ -873,23 +931,41 @@
       const wrapper = document.createElement('div');
       wrapper.className = 'code-block-wrapper';
 
-      // Create header with language and copy button
+      // Create header with language and action buttons
       const header = document.createElement('div');
       header.className = 'code-block-header';
+      
+      // Check if code is long enough to warrant an expand button
+      const codeText = codeBlock.textContent || '';
+      const isLongCode = codeText.split('\n').length > 15 || codeText.length > 500;
+      
       header.innerHTML = `
         <span class="code-block-language">${language}</span>
-        <button class="copy-code-btn" data-code-index="${index}">
-          <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
-            <rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect>
-            <path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path>
-          </svg>
-          Copy
-        </button>
+        <div class="code-block-actions">
+          ${isLongCode ? `
+            <button class="expand-code-btn" data-code-index="${index}">
+              <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
+                <path d="M8 3H5a2 2 0 0 0-2 2v3m18 0V5a2 2 0 0 0-2-2h-3m0 18h3a2 2 0 0 0 2-2v-3M3 16v3a2 2 0 0 0 2 2h3"></path>
+              </svg>
+              Expand
+            </button>
+          ` : ''}
+          <button class="copy-code-btn" data-code-index="${index}">
+            <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
+              <rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect>
+              <path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path>
+            </svg>
+            Copy
+          </button>
+        </div>
       `;
 
       // Create content wrapper and move the original <pre><code> inside (preserves highlighting)
       const content = document.createElement('div');
       content.className = 'code-block-content';
+      if (isLongCode) {
+        content.classList.add('collapsed');
+      }
       pre.dataset.sbWrapped = '1';
       content.appendChild(pre);
 
@@ -916,12 +992,42 @@
       // If messageDiv contains multiple elements, just append wrapper now
       messageDiv.appendChild(wrapper);
 
-      // Add click handler for copy button
+      // Add click handlers for copy and expand buttons
       const copyBtn = wrapper.querySelector('.copy-code-btn');
-      copyBtn.addEventListener('click', () => copyCodeToClipboard(codeBlock.textContent, copyBtn));
+      const expandBtn = wrapper.querySelector('.expand-code-btn');
+      
+      if (copyBtn) {
+        copyBtn.addEventListener('click', () => copyCodeToClipboard(codeBlock.textContent, copyBtn));
+      }
+      
+      if (expandBtn) {
+        expandBtn.addEventListener('click', () => toggleCodeExpansion(content, expandBtn));
+      }
     });
   }
 
+  function toggleCodeExpansion(content, button) {
+    const isCollapsed = content.classList.contains('collapsed');
+    
+    if (isCollapsed) {
+      content.classList.remove('collapsed');
+      button.innerHTML = `
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
+          <path d="M8 3H5a2 2 0 0 0-2 2v3m18 0V5a2 2 0 0 0-2-2h-3m0 18h3a2 2 0 0 0 2-2v-3M3 16v3a2 2 0 0 0 2 2h3"></path>
+        </svg>
+        Collapse
+      `;
+    } else {
+      content.classList.add('collapsed');
+      button.innerHTML = `
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
+          <path d="M8 3H5a2 2 0 0 0-2 2v3m18 0V5a2 2 0 0 0-2-2h-3m0 18h3a2 2 0 0 0 2-2v-3M3 16v3a2 2 0 0 0 2 2h3"></path>
+        </svg>
+        Expand
+      `;
+    }
+  }
+
   function copyCodeToClipboard(code, button) {
     navigator.clipboard.writeText(code).then(() => {
       const originalText = button.innerHTML;

static/styles.css

@@ -1074,57 +1074,82 @@
   overflow: auto;
 }
 
-/* Code blocks with copy button */
+/* Enhanced IDE-like Code blocks */
 .code-block-wrapper {
   position: relative;
   margin: 20px 0;
   border-radius: 12px;
   overflow: hidden;
   border: 1px solid var(--border);
-  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+  box-shadow: 0 8px 32px rgba(0, 0, 0, 0.12);
   background: var(--card);
+  transition: all 0.3s ease;
+}
+
+.code-block-wrapper:hover {
+  box-shadow: 0 12px 40px rgba(0, 0, 0, 0.15);
+  transform: translateY(-2px);
 }
 
 .code-block-header {
   display: flex;
   justify-content: space-between;
   align-items: center;
-  padding: 12px 16px;
-  background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--card) 100%);
-  border-bottom: 1px solid var(--border);
+  padding: 16px 20px;
+  background: linear-gradient(135deg, #2d2d2d 0%, #1e1e1e 100%);
+  border-bottom: 1px solid #404040;
   font-size: 12px;
-  color: var(--muted);
+  color: #a0a0a0;
   font-weight: 500;
+  position: relative;
+}
+
+.code-block-header::before {
+  content: '';
+  position: absolute;
+  top: 0;
+  left: 0;
+  right: 0;
+  height: 1px;
+  background: linear-gradient(90deg, transparent, var(--accent), transparent);
 }
 
 .code-block-language {
   font-weight: 700;
   text-transform: uppercase;
-  letter-spacing: 0.5px;
+  letter-spacing: 0.8px;
   color: var(--accent);
   display: flex;
   align-items: center;
-  gap: 6px;
+  gap: 8px;
+  font-size: 11px;
 }
 
 .code-block-language::before {
   content: '';
-  width: 8px;
-  height: 8px;
+  width: 10px;
+  height: 10px;
   border-radius: 50%;
   background: var(--accent);
   display: inline-block;
+  box-shadow: 0 0 8px rgba(var(--accent-rgb), 0.3);
 }
 
-.copy-code-btn {
+.code-block-actions {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+}
+
+.copy-code-btn, .expand-code-btn {
   display: flex;
   align-items: center;
   gap: 6px;
-  padding: 6px 12px;
-  background: var(--card);
-  border: 1px solid var(--border);
+  padding: 8px 14px;
+  background: rgba(255, 255, 255, 0.05);
+  border: 1px solid rgba(255, 255, 255, 0.1);
   border-radius: 8px;
-  color: var(--text-secondary);
+  color: #a0a0a0;
   font-size: 11px;
   font-weight: 500;
   cursor: pointer;
@@ -1133,7 +1158,7 @@
   overflow: hidden;
 }
 
-.copy-code-btn::before {
+.copy-code-btn::before, .expand-code-btn::before {
   content: '';
   position: absolute;
   top: 0;
@@ -1144,15 +1169,15 @@
   transition: left 0.5s ease;
 }
 
-.copy-code-btn:hover {
-  background: var(--card-hover);
+.copy-code-btn:hover, .expand-code-btn:hover {
+  background: rgba(255, 255, 255, 0.1);
   border-color: var(--accent);
-  color: var(--text);
+  color: #ffffff;
   transform: translateY(-1px);
-  box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
+  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2);
 }
 
-.copy-code-btn:hover::before {
+.copy-code-btn:hover::before, .expand-code-btn:hover::before {
   left: 100%;
 }
 
@@ -1163,25 +1188,26 @@
   transform: scale(1.05);
 }
 
-.copy-code-btn svg {
+.copy-code-btn svg, .expand-code-btn svg {
   width: 14px;
   height: 14px;
   transition: transform 0.2s ease;
 }
 
-.copy-code-btn:hover svg {
+.copy-code-btn:hover svg, .expand-code-btn:hover svg {
   transform: scale(1.1);
 }
 
 .code-block-content {
-  background: linear-gradient(135deg, #1a1a1a 0%, #1e1e1e 100%);
-  color: #e6e6e6;
-  padding: 20px;
+  background: linear-gradient(135deg, #0d1117 0%, #161b22 100%);
+  color: #e6edf3;
+  padding: 24px;
   overflow-x: auto;
-  font-family: 'JetBrains Mono', 'Fira Code', 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+  font-family: 'JetBrains Mono', 'Fira Code', 'Cascadia Code', 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
   font-size: 14px;
-  line-height: 1.6;
+  line-height: 1.7;
   position: relative;
+  border-left: 4px solid var(--accent);
 }
 
 .code-block-content::before {
@@ -1194,6 +1220,36 @@
   background: linear-gradient(90deg, transparent, var(--accent), transparent);
 }
 
+.code-block-content::after {
+  content: '';
+  position: absolute;
+  top: 0;
+  right: 0;
+  width: 1px;
+  height: 100%;
+  background: linear-gradient(180deg, transparent, var(--accent), transparent);
+}
+
+/* Line numbers for code blocks */
+.code-block-content.with-line-numbers {
+  counter-reset: line-number;
+}
+
+.code-block-content.with-line-numbers pre {
+  counter-increment: line-number;
+}
+
+.code-block-content.with-line-numbers pre::before {
+  content: counter(line-number);
+  position: absolute;
+  left: -40px;
+  width: 30px;
+  text-align: right;
+  color: #6e7681;
+  font-size: 12px;
+  user-select: none;
+}
+
 /* Enhanced syntax highlighting for code blocks */
 .code-block-content .hljs {
   background: transparent !important;
@@ -1220,12 +1276,75 @@
 
 .code-block-content .hljs-function {
   color: #d2a8ff;
+  font-weight: 500;
 }
 
 .code-block-content .hljs-variable {
   color: #ffa657;
 }
 
+.code-block-content .hljs-built_in {
+  color: #ffa657;
+  font-weight: 500;
+}
+
+.code-block-content .hljs-type {
+  color: #d2a8ff;
+  font-weight: 500;
+}
+
+.code-block-content .hljs-class {
+  color: #d2a8ff;
+  font-weight: 500;
+}
+
+.code-block-content .hljs-title {
+  color: #d2a8ff;
+  font-weight: 500;
+}
+
+.code-block-content .hljs-params {
+  color: #e6edf3;
+}
+
+.code-block-content .hljs-attr {
+  color: #79c0ff;
+}
+
+.code-block-content .hljs-tag {
+  color: #7ee787;
+}
+
+.code-block-content .hljs-name {
+  color: #7ee787;
+}
+
+.code-block-content .hljs-attribute {
+  color: #79c0ff;
+}
+
+.code-block-content .hljs-value {
+  color: #a5d6ff;
+}
+
+/* Code folding */
+.code-block-content.collapsed {
+  max-height: 200px;
+  overflow: hidden;
+  position: relative;
+}
+
+.code-block-content.collapsed::after {
+  content: '';
+  position: absolute;
+  bottom: 0;
+  left: 0;
+  right: 0;
+  height: 40px;
+  background: linear-gradient(transparent, #0d1117);
+  pointer-events: none;
+}
+
 /* Scrollbar styling for code blocks */
 .code-block-content::-webkit-scrollbar {
   height: 8px;

test_improvements.py

@@ -1,220 +0,0 @@
-#!/usr/bin/env python3
-"""
-Test script to verify code styling and heading numbering improvements
-"""
-import asyncio
-import sys
-import os
-
-# Add the project root to the Python path
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from utils.service.pdf import _apply_syntax_highlight
-
-def test_code_styling_improvements():
-    """Test that the code styling improvements work correctly"""
-    print("Testing code styling improvements...")
-    
-    # Test Python code highlighting
-    python_code = '''def calculate_fibonacci(n):
-    """Calculate the nth Fibonacci number"""
-    if n <= 1:
-        return n
-    return calculate_fibonacci(n-1) + calculate_fibonacci(n-2)
-
-# Test the function
-result = calculate_fibonacci(10)
-print(f"Fibonacci(10) = {result}")'''
-    
-    try:
-        # Escape the code first
-        escaped_code = python_code.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
-        result = _apply_syntax_highlight(escaped_code, 'python')
-        
-        print("✓ Python syntax highlighting test passed")
-        print("  Sample output (first 200 chars):")
-        print(f"  {result[:200]}...")
-        
-        # Check for key highlighting elements
-        assert '<font color=' in result, "Should have color highlighting"
-        assert 'def' in result, "Should highlight keywords"
-        assert 'print' in result, "Should highlight built-ins"
-        assert '#' in result, "Should highlight comments"
-        
-        return True
-    except Exception as e:
-        print(f"✗ Python syntax highlighting test failed: {e}")
-        return False
-
-def test_javascript_highlighting():
-    """Test JavaScript highlighting"""
-    print("\nTesting JavaScript highlighting...")
-    
-    js_code = '''function processData(data) {
-    // Process the input data
-    const result = data.map(item => {
-        return {
-            id: item.id,
-            name: item.name.toUpperCase(),
-            processed: true
-        };
-    });
-    
-    return result;
-}
-
-// Usage
-const data = [{id: 1, name: "test"}];
-const processed = processData(data);
-console.log(processed);'''
-    
-    try:
-        escaped_code = js_code.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
-        result = _apply_syntax_highlight(escaped_code, 'javascript')
-        
-        print("✓ JavaScript syntax highlighting test passed")
-        print("  Sample output (first 200 chars):")
-        print(f"  {result[:200]}...")
-        
-        # Check for key highlighting elements
-        assert '<font color=' in result, "Should have color highlighting"
-        assert 'function' in result, "Should highlight keywords"
-        assert 'console' in result, "Should highlight built-ins"
-        
-        return True
-    except Exception as e:
-        print(f"✗ JavaScript syntax highlighting test failed: {e}")
-        return False
-
-def test_json_highlighting():
-    """Test JSON highlighting"""
-    print("\nTesting JSON highlighting...")
-    
-    json_code = '''{
-    "name": "test",
-    "value": 123,
-    "active": true,
-    "items": null,
-    "data": {
-        "nested": "value"
-    }
-}'''
-    
-    try:
-        escaped_code = json_code.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
-        result = _apply_syntax_highlight(escaped_code, 'json')
-        
-        print("✓ JSON syntax highlighting test passed")
-        print("  Sample output (first 200 chars):")
-        print(f"  {result[:200]}...")
-        
-        # Check for key highlighting elements
-        assert '<font color=' in result, "Should have color highlighting"
-        assert 'true' in result, "Should highlight boolean values"
-        assert 'null' in result, "Should highlight null values"
-        
-        return True
-    except Exception as e:
-        print(f"✗ JSON syntax highlighting test failed: {e}")
-        return False
-
-def test_heading_extraction():
-    """Test heading extraction logic"""
-    print("\nTesting heading extraction...")
-    
-    sample_report = """# Introduction
-This is the introduction section.
-
-## Background
-Some background information.
-
-### Historical Context
-Historical details here.
-
-## Methodology
-Our approach to the problem.
-
-### Data Collection
-How we collected data.
-
-### Analysis
-How we analyzed the data.
-
-## Results
-The findings.
-
-## Conclusion
-Final thoughts."""
-    
-    import re
-    
-    # Extract headings using the same logic as the fix_heading_numbering function
-    heading_pattern = r'^(#{1,6})\s*(.*)$'
-    headings = []
-    lines = sample_report.split('\n')
-    
-    for i, line in enumerate(lines):
-        match = re.match(heading_pattern, line.strip())
-        if match:
-            level = len(match.group(1))
-            text = match.group(2).strip()
-            # Remove existing numbering if present
-            text = re.sub(r'^\d+\.?\s*', '', text)
-            headings.append({
-                'line_number': i,
-                'level': level,
-                'text': text,
-                'original_line': line
-            })
-    
-    print(f"✓ Extracted {len(headings)} headings:")
-    for h in headings:
-        print(f"  Level {h['level']}: {h['text']}")
-    
-    # Verify we found the expected headings
-    expected_headings = ['Introduction', 'Background', 'Historical Context', 'Methodology', 'Data Collection', 'Analysis', 'Results', 'Conclusion']
-    found_texts = [h['text'] for h in headings]
-    
-    for expected in expected_headings:
-        if expected in found_texts:
-            print(f"  ✓ Found: {expected}")
-        else:
-            print(f"  ✗ Missing: {expected}")
-    
-    return len(headings) >= len(expected_headings)
-
-async def main():
-    """Run all tests"""
-    print("Running code styling and heading numbering improvements tests...\n")
-    
-    # Test code styling improvements
-    python_test = test_code_styling_improvements()
-    js_test = test_javascript_highlighting()
-    json_test = test_json_highlighting()
-    
-    # Test heading extraction
-    heading_test = test_heading_extraction()
-    
-    print("\n" + "="*60)
-    print("TEST RESULTS:")
-    print(f"✓ Python highlighting: {'PASSED' if python_test else 'FAILED'}")
-    print(f"✓ JavaScript highlighting: {'PASSED' if js_test else 'FAILED'}")
-    print(f"✓ JSON highlighting: {'PASSED' if json_test else 'FAILED'}")
-    print(f"✓ Heading extraction: {'PASSED' if heading_test else 'FAILED'}")
-    
-    all_passed = python_test and js_test and json_test and heading_test
-    
-    if all_passed:
-        print("\n🎉 All tests passed!")
-        print("✓ Code styling improvements are working correctly")
-        print("✓ Heading extraction logic is working correctly")
-        print("✓ The PDF generation should now have better code styling")
-        print("✓ Reports should have properly numbered headings")
-    else:
-        print("\n❌ Some tests failed. Please check the output above.")
-    
-    print("\nNote: The heading re-numbering will use AI models during actual report generation.")
-    print("The extraction logic is tested here, but the AI re-numbering requires a live model.")
-
-if __name__ == "__main__":
-    asyncio.run(main())

utils/service/pdf.py

@@ -13,7 +13,7 @@ from utils.logger import get_logger
 logger = get_logger("PDF", __name__)
 
 
-def _parse_markdown_content(content: str, heading1_style, heading2_style, heading3_style, normal_style, code_style):
+async def _parse_markdown_content(content: str, heading1_style, heading2_style, heading3_style, normal_style, code_style):
     """
     Enhanced markdown parser that properly handles bold/italic formatting
     """
@@ -63,11 +63,15 @@ def _parse_markdown_content(content: str, heading1_style, heading2_style, headin
                 i += 1
             
             if code_lines:
-                # Mermaid diagrams → render via Kroki PNG for PDF
+                # Mermaid diagrams → render via Kroki PNG for PDF with retry logic
                 if language.lower() == 'mermaid':
                     try:
                         from reportlab.platypus import Image, Spacer
-                        img_bytes = _render_mermaid_png('\n'.join(code_lines))
+                        mermaid_code = '\n'.join(code_lines)
+                        # Use retry logic from diagram.py
+                        from helpers.diagram import _render_mermaid_with_retry
+                        img_bytes = await _render_mermaid_with_retry(mermaid_code)
+                        
                         if img_bytes and len(img_bytes) > 0:
                             import io
                             img = Image(io.BytesIO(img_bytes))
@@ -82,9 +86,9 @@ def _parse_markdown_content(content: str, heading1_style, heading2_style, headin
                             i += 1
                             continue
                         else:
-                            logger.warning("[PDF] Mermaid render returned empty image, falling back to code block")
+                            logger.warning("[PDF] Mermaid render returned empty image after retries, falling back to code block")
                     except Exception as me:
-                        logger.warning(f"[PDF] Mermaid render failed, falling back to code block: {me}")
+                        logger.warning(f"[PDF] Mermaid render failed after retries, falling back to code block: {me}")
                     
                     # Fallback: render as code block with mermaid syntax
                     from reportlab.platypus import XPreformatted, Paragraph
@@ -828,12 +832,12 @@ async def generate_report_pdf(report_content: str, user_id: str, project_id: str
         story = []
         
         # Add title
-        story.append(Paragraph("Study Report", title_style))
+        story.append(Paragraph("StudyBuddy Report", title_style))
         story.append(Paragraph(f"<i>Generated on {datetime.now().strftime('%B %d, %Y at %I:%M %p')}</i>", normal_style))
         story.append(Spacer(1, 20))
         
         # Enhanced markdown parser with proper formatting
-        story.extend(_parse_markdown_content(report_content, heading1_style, heading2_style, heading3_style, normal_style, code_style))
+        story.extend(await _parse_markdown_content(report_content, heading1_style, heading2_style, heading3_style, normal_style, code_style))
         
         # Add references section if sources provided
         if sources: